Conditionally sum broadcasted dimensions in Elemwise and BroadcastTo gradients

Author: brandonwillard

aesara/tensor/elemwise.py

@@ -543,7 +543,6 @@ def connection_pattern(self, node):
         return [[True for output in node.outputs] for ipt in node.inputs]
 
     def L_op(self, inputs, outs, ograds):
-        from aesara.tensor.math import sum as at_sum
 
         # Compute grad with respect to broadcasted input
         rval = self._bgrad(inputs, outs, ograds)
@@ -574,18 +573,9 @@ def L_op(self, inputs, outs, ograds):
             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
+            rval[i] = aesara.tensor.extra_ops.sum_broadcasted_dims(
+                rval[i], ipt, outs[0].type.shape
+            )
 
         return rval
 

aesara/tensor/extra_ops.py

@@ -1,6 +1,6 @@
 from collections.abc import Collection
 from functools import reduce
-from typing import Iterable, Set, Tuple, Union
+from typing import Iterable, Optional, Sequence, Set, Tuple, Union
 
 import numpy as np
 import numpy.core.numeric
@@ -1669,19 +1669,11 @@ def grad(self, inputs, outputs_gradients):
 
         d_wrt_a = broadcast_to(dout, shape).sum(axis=new_dims)
 
-        # Determine the dimensions that were broadcast
-        _, static_shape = at.infer_static_shape(shape)
-
-        # TODO: This needs to be performed at run-time when static shape
-        # information isn't available.
-        bcast_sums = [
-            i
-            for i, (a_s, s_s) in enumerate(zip(a.type.shape, static_shape[-a.ndim :]))
-            if a_s == 1 and s_s != 1
-        ]
-
-        if bcast_sums:
-            d_wrt_a = d_wrt_a.sum(axis=bcast_sums, keepdims=True)
+        # Determine the dimensions that were broadcast and sum them
+        static_out_shape = tuple(
+            s.data if isinstance(s, Constant) else None for s in shape[-a.ndim :]
+        )
+        d_wrt_a = sum_broadcasted_dims(d_wrt_a, a, static_out_shape)
 
         return [d_wrt_a] + [
             grad_undefined(self, i, shp) for i, shp in enumerate(shape, 1)
@@ -1808,6 +1800,46 @@ def broadcast_arrays(*args: TensorVariable) -> Tuple[TensorVariable, ...]:
     return tuple(broadcast_to(a, broadcast_shape(*args)) for a in args)
 
 
+def sum_broadcasted_dims(
+    value: TensorVariable,
+    inp: TensorVariable,
+    out_shape: Sequence[Optional[int]],
+) -> TensorVariable:
+    """Sum dimensions in `value` that are broadcasted between `inp`'s shape and `out_shape`.
+
+    For ambiguous cases, this builds a graph that determine whether or not
+    dimensions are to be summed at run-time.
+
+    """
+    dims_to_sum = ()
+    ambiguous_dim_conds = ()
+
+    in_shape = inp.type.shape
+
+    for i, (s1, s2) in enumerate(zip(in_shape, out_shape)):
+        if s1 == 1 and s2 != 1:
+            dims_to_sum += (i,)
+        elif s1 is None and s2 != 1:
+            ambiguous_dim_conds += (
+                (i, aes.eq(at.scalar_from_tensor(inp.shape[i]), 1)),
+            )
+
+    if dims_to_sum:
+        value = at_sum(value, axis=dims_to_sum, keepdims=True)
+
+    if ambiguous_dim_conds:
+        from aesara.ifelse import ifelse
+
+        for i, cond in ambiguous_dim_conds:
+            value = ifelse(
+                cond,
+                at_sum(value, axis=i, keepdims=True),
+                value,
+            )
+
+    return value
+
+
 __all__ = [
     "searchsorted",
     "cumsum",

tests/tensor/test_elemwise.py

@@ -8,8 +8,9 @@
 
 import aesara
 import aesara.scalar as aes
+import aesara.tensor as at
 import tests.unittest_tools as utt
-from aesara.compile.mode import Mode
+from aesara.compile.mode import Mode, get_default_mode
 from aesara.configdefaults import config
 from aesara.graph.basic import Apply, Variable
 from aesara.graph.fg import FunctionGraph
@@ -889,6 +890,39 @@ def test_invalid_static_shape(self):
         ):
             x + y
 
+    def test_grad_sum_bcast_input_dims(self):
+        """Make sure broadcasted dimensions in the gradients are summed when static shape information isn't available."""
+        Y = matrix("Y")
+        X = matrix("X")
+        X_grad = aesara.grad((X + Y).sum(), wrt=X)
+
+        mode = get_default_mode().including("fast_run")
+
+        X_grad_fn = aesara.function([X, Y], X_grad, mode=mode)
+        res = X_grad_fn(np.ones((1, 5)), np.ones((5, 5)))
+        assert np.array_equal(res, np.array([[5.0, 5.0, 5.0, 5.0, 5.0]]))
+
+        # When the shapes are known at compile-time, the compiled graph should
+        # simplify
+        Y = tensor(np.float64, shape=(5, None), name="Y")
+        X = tensor(np.float64, shape=(1, 5), name="X")
+        X_grad = aesara.grad((X + Y).sum(), wrt=X)
+
+        X_grad_fn = aesara.function([X, Y], X_grad, mode=mode)
+        res = X_grad_fn(np.ones((1, 5)), np.ones((5, 5)))
+        assert np.array_equal(res, np.array([[5.0, 5.0, 5.0, 5.0, 5.0]]))
+
+        assert X_grad_fn.maker.fgraph.apply_nodes
+
+    def test_grad_of_grad(self):
+        """This tests a special case in which the static shapes of a `DimShuffle` and its gradient don't match."""
+        a = at.vector("a")
+
+        out = aesara.grad((a * a).sum(), a).sum()
+        out = aesara.grad(out, a)
+
+        assert out.type.shape == (None,)
+
 
 def test_not_implemented_elemwise_grad():
     # Regression test for unimplemented gradient in an Elemwise Op.

tests/tensor/test_extra_ops.py

@@ -1318,6 +1318,7 @@ def test_memory_leak(self):
         [
             [lambda x: broadcast_to(x, (1,)), (1,)],
             [lambda x: broadcast_to(x, (6, 2, 5, 3)), (1,)],
+            [lambda x: broadcast_to(x, (6, 2, 5, 3)), (1,)],
             [lambda x: broadcast_to(x, (6, 2, 5, 3)), (5, 1)],
             [lambda x: broadcast_to(x, (6, 2, 1, 3)), (2, 1, 3)],
         ],