Guarantee that branch shapes match in the gradient of IfElse

Author: brandonwillard

aesara/ifelse.py

@@ -24,7 +24,7 @@
 from aesara.graph.op import _NoPythonOp
 from aesara.graph.rewriting.basic import GraphRewriter, in2out, node_rewriter
 from aesara.graph.type import HasDataType, HasShape
-from aesara.tensor.shape import Reshape, Shape, SpecifyShape, Unbroadcast
+from aesara.tensor.shape import Reshape, Shape, SpecifyShape, Unbroadcast, specify_shape
 
 
 if TYPE_CHECKING:
@@ -254,33 +254,42 @@ def grad(self, ins, grads):
         # Since input true/false entries must have the same dtypes, we need to
         # cast the zeros to the corresponding `grads` dtypes and not the input
         # dtypes.
-        inputs_true_grad = (
-            [condition]
-            + grads
-            + [
-                at.basic.zeros_like(t, dtype=grads[i].dtype)
-                for i, t in enumerate(inputs_true_branch)
-            ]
+        # The `grads` can also have different shapes than the `inputs`, so we
+        # effectively assert that the shapes are preserved in each branch.
+        # TODO FIXME: This doesn't seem like a sufficient solution to the
+        # problem.
+        inputs_true_grads = if_true_op(
+            *(
+                [condition]
+                + [specify_shape(g, i.shape) for g, i in zip(grads, inputs_true_branch)]
+                + [
+                    at.basic.zeros_like(t, dtype=grads[i].dtype)
+                    for i, t in enumerate(inputs_true_branch)
+                ]
+            ),
+            return_list=True,
         )
-        inputs_false_grad = (
-            [condition]
-            + [
-                at.basic.zeros_like(f, dtype=grads[i].dtype)
-                for i, f in enumerate(inputs_false_branch)
-            ]
-            + grads
+        inputs_false_grads = if_false_op(
+            *(
+                [condition]
+                + [
+                    at.basic.zeros_like(f, dtype=grads[i].dtype)
+                    for i, f in enumerate(inputs_false_branch)
+                ]
+                + [
+                    specify_shape(g, i.shape)
+                    for g, i in zip(grads, inputs_false_branch)
+                ]
+            ),
+            return_list=True,
         )
 
         # `condition` does affect the elements of the output so it is connected.
         # For the sake of making the gradient convenient we assume that
         # condition + epsilon always triggers the same branch as condition
         condition_grad = condition.zeros_like().astype(config.floatX)
 
-        return (
-            [condition_grad]
-            + if_true_op(*inputs_true_grad, return_list=True)
-            + if_false_op(*inputs_false_grad, return_list=True)
-        )
+        return [condition_grad] + inputs_true_grads + inputs_false_grads
 
     def make_thunk(self, node, storage_map, compute_map, no_recycling, impl=None):
         cond = node.inputs[0]

tests/test_ifelse.py

@@ -718,3 +718,56 @@ def test_nested():
     linker = aesara.link.vm.VMLinker(lazy=True)
     f = function([c1, c2, x1, x2], t4, mode=Mode(linker=linker, optimizer="fast_run"))
     assert f(1, 0, np.array(10, dtype=x1.dtype), 0) == 20.5
+
+
+def test_DimShuffle_drop():
+    c = scalar("c")
+    x = scalar("x")
+    y = vector("y")
+
+    cost = ifelse(c, x.dimshuffle("x"), y).sum()
+
+    # Sum{acc_dtype=float64} [id A] <TensorType(float64, ())>
+    #  |if{} [id B] <TensorType(float64, (None,))>
+    #    |c [id C] <TensorType(float64, ())>
+    #    |InplaceDimShuffle{x} [id D] <TensorType(float64, (1,))>
+    #    | |x [id E] <TensorType(float64, ())>
+    #    |y [id F] <TensorType(float64, (None,))>
+
+    out = aesara.grad(cost, y)
+    assert out.type.shape == (None,)
+
+    out = aesara.grad(cost, x)
+
+    #
+    # `DimShuffle.L_op` `inputs`
+    #
+    # x [id A] <TensorType(float64, ())>
+
+    #
+    # `DimShuffle.L_op` `outputs`
+    #
+    # InplaceDimShuffle{x} [id B] <TensorType(float64, (1,))>
+    #  |x [id A] <TensorType(float64, ())>
+
+    #
+    # `DimShuffle.L_op` `output_grads`
+    #
+    # if{} [id C] <TensorType(float64, (None,))>
+    #  |c [id D] <TensorType(float64, ())>
+    #  |Elemwise{second} [id E] <TensorType(float64, (None,))>
+    #  | |if{} [id F] <TensorType(float64, (None,))>
+    #  | | |c [id D] <TensorType(float64, ())>
+    #  | | |InplaceDimShuffle{x} [id B] <TensorType(float64, (1,))>
+    #  | | |y [id G] <TensorType(float64, (None,))>
+    #  | |InplaceDimShuffle{x} [id H] <TensorType(float64, (1,))>
+    #  |   |Elemwise{second,no_inplace} [id I] <TensorType(float64, ())>
+    #  |     |Sum{acc_dtype=float64} [id J] <TensorType(float64, ())>
+    #  |     | |if{} [id F] <TensorType(float64, (None,))>
+    #  |     |TensorConstant{1.0} [id K] <TensorType(float64, ())>
+    #  |Elemwise{second,no_inplace} [id L] <TensorType(float64, (1,))>
+    #    |InplaceDimShuffle{x} [id B] <TensorType(float64, (1,))>
+    #    |InplaceDimShuffle{x} [id M] <TensorType(float64, (1,))>
+    #      |TensorConstant{0.0} [id N] <TensorType(float64, ())>
+
+    assert out.type.shape == ()