Replace uses of Rebroadcast by SpecifyShape

[ricardoV94]

This one is not as straightforward as I was hoping, mainly because of the current dependency on the ability to "unbroadcast" variables (i.e., to mask known broadcastable dimensions). This was done via unbroadcast and patternbroadcast which combined both unbroadcast and addbroadcast

"Unbroadcasting" is still used explicitly in Scan and broadcast_like.

Changes

For now explicit "unbroadcasting" goes through the more limited MaskBroadcastable Op, adapted from the old Rebroadcast. Everything else that used to rely on Rebroadcast uses SpecifyShape or the new thin wrapper specify_broadcastable.

Addresses part of #748 and #917

Closes #955

[brandonwillard]

Ideally we wouldn't need this Op at all but I fear there are some deeply-ingrained dependencies on this...

Sounds like this is the core of our work: finding these dependencies and addressing them directly.

[ricardoV94]

The last places that seem to make explicit use of "unbroadcasting" are these:

aesara/aesara/tensor/basic_opt.py

Lines 156 to 165 in 6cca25e

	# the template may have 1s in its shape without being broadcastable
	if rval.broadcastable != template.broadcastable:
	rval = unbroadcast(
	rval,
	*[
	i
	for i in range(rval.ndim)
	if rval.broadcastable[i] and not template.broadcastable[i]
	],
	)

aesara/aesara/scan/basic.py

Lines 750 to 758 in 6cca25e

	# We need now to allocate space for storing the output and copy
	# the initial state over. We do this using the expand function
	# defined in scan utils
	sit_sot_scan_inputs.append(
	expand_empty(
	at.unbroadcast(shape_padleft(actual_arg), 0),
	actual_n_steps,
	)
	)

aesara/aesara/scan/basic.py

Lines 878 to 885 in 6cca25e

	# we need to see if we need to pad our sequences with an
	# unbroadcastable dimension; case example : we return an
	# output for which we want all intermediate. If n_steps is 1
	# then, if we return the output as given by the innner function
	# this will represent only a slice and it will have one
	# dimension less.
	if isinstance(inner_out.type, TensorType) and return_steps.get(pos, 0) != 1:
	outputs[pos] = at.unbroadcast(shape_padleft(inner_out), 0)

aesara/aesara/scan/basic.py

Lines 1012 to 1017 in 6cca25e

	sit_sot_scan_inputs.append(
	expand_empty(
	at.unbroadcast(shape_padleft(input.variable), 0),
	actual_n_steps,
	)
	)

[ricardoV94]

The issue here seems to be creating an EmptyAlloc that does not have a fixed shape of 1 on the leftmost dimension. I didn't dig to find out why/whether this is actually needed

[codecov]

Codecov Report

Merging #915 (4cfaaa5) into main (ccfe2d3) will decrease coverage by 0.02%.
The diff coverage is 91.44%.

❗ Current head 4cfaaa5 differs from pull request most recent head 8a6a871. Consider uploading reports for the commit 8a6a871 to get more accurate results

@@            Coverage Diff             @@
##             main     #915      +/-   ##
==========================================
- Coverage   79.27%   79.24%   -0.03%     
==========================================
  Files         152      152              
  Lines       47965    47848     -117     
  Branches    10923    10902      -21     
==========================================
- Hits        38023    37919     -104     
+ Misses       7438     7423      -15     
- Partials     2504     2506       +2     

Impacted Files	Coverage Δ
aesara/compile/function/pfunc.py	83.41% <ø> (ø)
aesara/scan/op.py	85.39% <ø> (-0.09%)	⬇️
aesara/tensor/blas.py	79.63% <33.33%> (-0.08%)	⬇️
aesara/tensor/nnet/conv.py	79.40% <33.33%> (-0.57%)	⬇️
aesara/tensor/nnet/opt.py	43.02% <66.66%> (+0.05%)	⬆️
aesara/tensor/nnet/batchnorm.py	76.92% <81.81%> (-0.12%)	⬇️
aesara/tensor/basic_opt.py	85.66% <83.33%> (-0.50%)	⬇️
aesara/ifelse.py	49.71% <100.00%> (-0.15%)	⬇️
aesara/link/jax/dispatch.py	80.20% <100.00%> (-1.69%)	⬇️
aesara/link/numba/dispatch/tensor_basic.py	100.00% <100.00%> (+2.06%)	⬆️
... and 36 more

[ricardoV94]

Tests are passing!

[ricardoV94]

I pushed a commit that removes the use of "unbroadcasting" in broadcast_like. This seems to have been introduced a long time ago in 9a45333, perhaps due to old rewrite limitations concerning change of static broadcastable type shape? It might no longer be relevant after #711, but unfortunately there were no specific tests targeting the functionality to be able to reason more directly.

In that same commit, several alloc calls were replaced by broadcast_like, which could mean that most of the remaining uses of broadcast_like might not be needed anymore? This would facilitate #288

If "unbroadcasting" is not necessary in broadcast_like, the only remaining use is in the creation of Scan. This is still clearly needed at the moment (even though it should not!) as indicated by the tests mentioned above.

[brandonwillard]

If "unbroadcasting" is not necessary in broadcast_like, the only remaining use is in the creation of Scan. This is still clearly needed at the moment (even though it should not!) as indicated by the tests mentioned above.

Let's remove it from Scan. I'll take a look today.

[ricardoV94]

If "unbroadcasting" is not necessary in broadcast_like, the only remaining use is in the creation of Scan. This is still clearly needed at the moment (even though it should not!) as indicated by the tests mentioned above.

Let's remove it from Scan. I'll take a look today.

All the tests passed without the use of unbroadcast in broadcast_like, so Scan is the only place where it is used now

[ricardoV94]

Renamed MaskBroadcastable to Unbroadcast

[brandonwillard]

This was/is a big improvement!