Wrong gradients when inputs are dynamically broadcasted

[ricardoV94]

This bug is an unexpected consequence of #928 and rewrites that make certain assumptions: #1089 (comment)

import aesara
import aesara.tensor as at
import numpy as np

x_row = at.row("x_row")
x_matrix = at.matrix("x_matrix")
y = at.matrix("y")

x_row_grad = at.grad(at.sum(x_row + y), wrt=x_row)
x_matrix_grad = at.grad(at.sum(x_matrix + y), wrt=x_matrix)

f_row = aesara.function([x_row, y], x_row_grad)
print(f_row(np.ones((1, 5)), np.ones((5, 5))))
# [[5. 5. 5. 5. 5.]]

f_matrix = aesara.function([x_matrix, y], x_matrix_grad)
print(f_matrix(np.ones((1, 5)), np.ones((5, 5))))
# [[1. 1. 1. 1. 1.]
#  [1. 1. 1. 1. 1.]
#  [1. 1. 1. 1. 1.]
#  [1. 1. 1. 1. 1.]
#  [1. 1. 1. 1. 1.]]

The faulty logic is found here:

aesara/aesara/tensor/elemwise.py

Lines 552 to 570 in 7f4e0ab

	# sum out the broadcasted dimensions
	for i, ipt in enumerate(inputs):
	if isinstance(rval[i].type, (NullType, DisconnectedType)):
	continue
	# List of all the dimensions that are broadcastable for input[i] so
	# we can sum over them
	# TODO: only count dimensions that were effectively broadcasted
	to_sum = [
	j
	for j, bcast in enumerate(ipt.type.broadcastable)
	if bcast and not outs[0].broadcastable[j]
	]
	if to_sum:
	sr = at_sum(rval[i], axis=to_sum, keepdims=True)
	rval[i] = sr
	return rval

This is also likely a problem in the grad of BroadcastTo which calls infer_broadcastable and which defaults to assuming something will not have broadcasted if a static shape of 1 can't be inferred.

aesara/aesara/tensor/extra_ops.py

Line 1593 in 7f8af9b

_, shape_bcast = at.infer_broadcastable(shape)

aesara/aesara/tensor/basic.py

Line 1313 in 7f8af9b

def infer_broadcastable(shape):

And also GEMM since #986

I am not sure if there's a good solution to this problem, as we would need an expression with different output shapes depending on whether the runtime inputs are broadcasted or not.

Solution might look something like: #1089 (comment)

[aseyboldt]

Uff, this looks bad.

This makes me wish we had dimension objects even more, I think that would solve this problem?
But I don't see how we could introduce that as a strict requirement without breaking backward compatibility...

If each dimension had an object associated, and we only allow broadcasting if those object are identical, we would always know when we create nodes how broadcasting would work, even before we know if a dimension happens to have length 1.

So something like:

dim0 = at.Dimension(name="dim0")
dim1 = at.Dimension(name="dim1")
dim2 = at.Dimension(name="dim2")

x = at.dvector("x", dims=(dim0,))
y = at.dvector("y", dims=(dim1,))

x + y  # dims == (dim0, dim1)

x = at.dvector("x", dims=(dim0,))
y = at.dvector("y", dims=(dim0,))

x + y  # dims=(dim0,)

X = at.dmatrix("X", dims=(dim0, dim1))
Y = at.dmatrix("Y", dims=(dim0, dim2))

X + Y  # dims == (dim0, dim1, dim2)

That way we'd always statically know the result shape/dims of broadcasting ops.

[brandonwillard]

This makes me wish we had dimension objects even more, I think that would solve this problem?

It looks like the dimension objects you're talking about would provide essentially the same information/accomplish the same things as specify_shape (alongside the existing shape inference). Perhaps the main difference is that these dimension objects are implicitly intended to work at the Type level? If so, the advantages of that approach aren't apparent at the moment.

[aseyboldt]

No, it is a pretty different idea than specify_shape. specify_shape just fixes the shape to known integer values, it provides no information if two different arrays that happen to have the same shape actually refer to the same dimension.

By dimension I mean some constant thing that has a clear identity and might for instance also have associated coordinates.

So if we have dimensions time and country, and x.dims == (time,); y.dims == (country,), those two axes could never broadcast together, even if by accident we had the same number of time points as countries, because their respective dimensions refer to different, incompatible things.

And if we have two arrays with the same dimension, we still wouldn't know the shape, but only that those arrays have the same, and a conceptually compatible dimension. A type like Array<dims=(ountry,)> would mean that we are providing one value for each country (whatever number of countries we have). It never makes sense to directly add this to Array<dims=('time')>, ie "one value for each point in time", even if n_time == n_countries. So for instance we can redefine the sum as a broadcasting sum with output Array<dims=('time', 'country')>.

This is essentially how broadcasting works in xarray, and in my opinion it is way cleaner, prevents lots of bugs and is in general a joy to work with, compared with the relatively messy numpy behavior, where for instance all hell can break lose if it just so happens that one of your dimensions has length 1. It does require a bit of acclimatization though. :-)

Edit
But I guess unless we can find a way to solve the compatibility problems, I'm not sure this is the right place to discuss this I guess.

[brandonwillard]

specify_shape just fixes the shape to known integer values, it provides no information if two different arrays that happen to have the same shape actually refer to the same dimension.

specify_shape can be used to assign the same non-constant Variable to arbitrary dimensions in the shapes of multiple Variables. Combine this with shape inference and it's possible to determine that distinct Variables share exactly the same shapes in those dimensions fixed by specify_shape.

import aesara
import aesara.tensor as at
from aesara.graph.fg import FunctionGraph
from aesara.graph.opt_utils import optimize_graph
from aesara.tensor.basic_opt import ShapeFeature


shape_dim_1 = at.lscalar("time")
shape_dim_2 = at.lscalar("country")

x = at.vector("x")
x = at.specify_shape(x, (shape_dim_1,))

y = at.matrix("y")
y = at.specify_shape(y, (shape_dim_1, shape_dim_2))


shape_feature = ShapeFeature()
fg = FunctionGraph(outputs=[x.shape[0], y.shape[0]], features=[shape_feature])

aesara.dprint(fg)
# Subtensor{int64} [id A] 2
#  |Shape [id B] 1
#  | |SpecifyShape [id C] 0
#  |   |x [id D]
#  |   |time [id E]
#  |ScalarConstant{0} [id F]
# Subtensor{int64} [id G] 5
#  |Shape [id H] 4
#  | |SpecifyShape [id I] 3
#  |   |y [id J]
#  |   |time [id E]
#  |   |country [id K]
#  |ScalarConstant{0} [id L]

fg_opt = optimize_graph(fg)

aesara.dprint(fg_opt)
# time [id A]
# time [id A]

shape_feature.shape_of
# {x: (Shape_i{0}.0,),
#  time: (),
#  SpecifyShape.0: (time,),
#  Shape.0: (TensorConstant{1},),
#  ScalarConstant{0}: (),
#  Subtensor{int64}.0: (),
#  y: (Shape_i{0}.0, Shape_i{1}.0),
#  country: (),
#  SpecifyShape.0: (time, country),
#  Shape.0: (TensorConstant{2},),
#  ScalarConstant{0}: (),
#  Subtensor{int64}.0: (),
#  MakeVector{dtype='int64'}.0: (TensorConstant{2},),
#  InplaceDimShuffle{}.0: (),
#  MakeVector{dtype='int64'}.0: (TensorConstant{1},)}

The main distinction I see here is the desire to refer to dimensions independently from their shapes, but that's not an entirely meaningful distinction in this symbolic context, because a symbolic indexed shape can serve as a direct proxy for a dimension.

[aseyboldt]

Oh, that's neat, I didn't realize you could put in variables! I'll have to play with this and see where I can get with it.

The different broadcasting behavior still does seem like a different issue though, doesn't it?

[brandonwillard]

Oh, that's neat, I didn't realize you could put in variables! I'll have to play with this and see where I can get with it.

There are a lot of places where we could/should be using shape inference, but we don't, so, while this stuff is possible, it's not necessarily used/as integrated as one would expect.

The different broadcasting behavior still does seem like a different issue though, doesn't it?

I'm not entirely clear on this situation yet; I only mentioned specify_shape and shape inference because, when combined, they should—theoretically—allow us to determine (non-)equivalent shapes when static/Type-level information isn't available.

[ricardoV94]

It seems like we might need a new Op that unbroadcasts (reduces) arrays to a given shape or leaves the input unchanged.

Check https://mostafa-samir.github.io/auto-diff-pt2/#unbroadcasting-adjoints

https://github.com/Mostafa-Samir/Hands-on-Intro-to-Auto-Diff/blob/29a9e5157421e2603846a15bceff21d3b2104f3d/autodiff/grads.py#L103

Something like:

x = at.matrix("x")

# at.reduce_to is probably a better name
# maybe sum is all we will ever need
y1 = at.unbroadcast_to(x, shape=(1, 5), reduce_op="sum”) 
y1.eval({x: np.ones((5, 5))})  # [[5, 5, 5, 5, 5]]

# This won't do anything, but shape may be only 
# known at runtime, as in the example in this issue!
y2 = at.unbroadcast_to(x, shape=(5, 5))
y2.eval({x: np.ones((5, 5))})  # np.ones((5, 5))

# If the shape is not compatible with something that could 
# have been broadcasted to the input shape, an error is raised
y3 = at.unbroadcast_to(x, shape=(2, 5))
y3.eval({x: np.ones((5, 5))})  # ValueError

This was also brought up by @Sayam753 and @purna135 in Slack in relation to their work on batched solve where dynamic unbroadcasting gradients also crops up. It was that discussion that led me to suspect of this bug!

Edit: This may be possible already without a specialized Op, if sum allows for symbolic axis? Does it?

In that case we could cook a helper pretty quickly, and perhaps add some rewrite in case the axis are constant folded during compilation/ and a sum with constant axis is more efficient.

Edit: Sum does not allow for variable axis

[ricardoV94]

Edit: mentioned other related issues in the top comment.

[brandonwillard]

Let's try to move to the Blockwise form of this problem/situation. That way, we can attempt to make some progress on #695 (e.g. finish implementing Blockwise.Lop in #757) and address these issues (via inheritance from the Elemwise case).

[brandonwillard]

This is also likely a problem in the grad of BroadcastTo which calls infer_broadcastable and which defaults to assuming something will not have broadcasted if a static shape of 1 can't be inferred.

Yeah, it looks like we might need to add symbolic conditions for those dimensions and let them be simplified later via shape inference and/or constant folding. This is similar to what we do when constructing broadcasted shape graphs.

[aseyboldt]

Maybe it would help if we added an op that precisely computes what actually happens when we broadcast several arrays?

So a wrapper around something like this (I really hope this can be simplified a bit, though). For known shapes this could be constant-folded.

def broadcast_patterns_numpy(*shapes):
    """Given input shapes predict broadcasting behavior.
    
    Returns
    -------
    - output_shape: List[int]
      The shape of the broadcasted array
    - did_broadcast: List[List[bool]]
      For each input and each output dimension: True, if the input
      was broadcasted in this particular axis
    - dim_index: List[List[Optional[int]]]
      For each input and each output dimension: An index that indicates
      which axis of the input array corresponds to the output axis. If
      the input did not have a axis, None
    """
    out_rank = max(len(shape) for shape in shapes)
    shapes_rev = [shape[::-1] for shape in shapes]

    # by_rank[i] contains the lengths of the
    # input dimensions in position i, counted
    # from the back
    by_rank = [[] for _ in range(out_rank)]
    for shape_rev in shapes_rev:
        i = -1
        for i, length in enumerate(shape_rev):
            by_rank[i].append(length)
        for j in range(i + 1, out_rank):
            by_rank[j].append(None)
    
    does_broadcast_rev = [[] for _ in shapes]
    dim_index_rev = [[] for _ in shapes]
    output_shape_rev = []
    for rank, lengths in enumerate(by_rank):
        lengths_non_zero = [length for length in lengths if length is not None and length != 1]
        if not lengths:
            out_length = 1
        else:
            out_length = lengths_non_zero[0]
        for length in lengths_non_zero[1:]:
            if length != out_length:
                raise ValueError(f"Could not broadcast at rev rank {rank}: {lengths}")

        output_shape_rev.append(length)

        for input, length in enumerate(lengths):
            if length is None:
                does_broadcast_rev[input].append(True)
                dim_index_rev[input].append(None)
            else:
                does_broadcast_rev[input].append(length == 1 and out_length != 1)
                n_before = sum(idx is not None for idx in dim_index_rev[input])
                dim_index_rev[input].append(len(shapes[input]) - n_before - 1)
        
    does_broadcast = [bools[::-1] for bools in does_broadcast_rev]
    dim_index = [idxs[::-1] for idxs in dim_index_rev]
    output_shape = output_shape_rev[::-1]
    
    return output_shape, does_broadcast, dim_index

For example if we broadcast arrays of shape (), (2, 5, 0) and (1, 5, 1) we'd get

out_shape, did_broadcast, dim_index = broadcast_patterns_numpy((), (2, 5, 0), (1, 5, 1))

out_shape == [2, 5, 0]
did_broadcast == [
    [True, True, True],   # The first argument was broadcast in all axes
    [False, False, False],  # The second argument wasn't broadcasted at all
    [True, False, True]   # The third argument was broadcasted in the first and last axis
]
dim_index == [
    [None, None, None],   # New axes where created for all output axis for the first arg
    [0, 1, 2],   # For the second arg all axes were used and their order didn't change
    [0, 1, 2]
]

The dim_index is a little redundant in numpy broadcasting, but it would be really helpful in xarray broadcasting. ;-)

[brandonwillard]

So a wrapper around something like this (I really hope this can be simplified a bit, though). For known shapes this could be constant-folded.

The function I linked to in my previous comment already computes broadcast shapes.

[ricardoV94]

For example if we broadcast arrays of shape (), (2, 5, 0) and (1, 5, 1) we'd get

What's up with that shape of (2, 5, 0)?

I think that combo is invalid as per numpy broadcasting rules due to the zero?

Maybe it would help if we added an op that precisely computes what actually happens when we broadcast several arrays?

That sounds valid, but it seems a more convoluted answer if you ONLY want to fix this problem.

In these cases we have the input and broadcasted gradient output, so the only thing we need is to reduce that gradient along the dimensions that where of size 1 in the input.

Actually, in the Elemwise case we don't even have to worry about new dims, because make_node adds Dimshuffles to the inputs to align the number of dims (but we will perhaps remove that one day)

So what we need is just:

# perform method of new Op boils down to this
def unbroadcast_to(x, shape):
  axis_to_sum = [
    i
    for i, (s, xs) in enumerate(zip(shape, x.shape))
    if s==1 and xs !=1
  ]

  if not axis_to_sum:
    return x
  return np.sum(x, axis=axis_to_sum, keepdims=True)

In the grad where this issue would crop up we would do something like

grad = unbroadcast_to(bcast_grad, shape=input.shape)

And then we could have some rewrites to try to get rid of this Op during compilation. For instance:

1. if the target shape and the input shapes are known to be equivalent, we can remove the Op

2. if the target shape has some (constant-folded) 1s we can replace the original input by one where we summed the known 1s dimensions already.

3. Hopefully there's already a rewrite to get rid of useless sum along dimensions of size 1. If not, we can add one.

4. if the target shape has no (constant-folded) 1s we can remove the Op (or replace by some Asserts related to the shape if we want to)