Repository: aesara-devs/aesara

Issue:
C Elemwise implementation doesn't broadcast variables
I'm seeing a very weird error in `Elemwise`:

First, here's a basic broadcasting operations in NumPy:

```python
import numpy as np

x = np.array([[-1.32720483],
              [ 0.23442016]])

m = np.array([0., 0.])

z = x - m
```

```python
>>> z
array([[-1.32720483, -1.32720483],
       [ 0.23442016,  0.23442016]])
```

In Aesara, here's the equivalent operation using `TensorConstant`s:

```python
import aesara
import aesara.tensor as at


x_at = at.as_tensor(x)
m_at = at.as_tensor(m)

z_at = x_at - m_at

```

```python
>>> aesara.dprint(z_at)
Elemwise{sub,no_inplace} [id A] ''
 |TensorConstant{[[-1.32720..23442016]]} [id B]
 |InplaceDimShuffle{x,0} [id C] ''
   |TensorConstant{(2,) of 0.0} [id D]
```

The resulting graph is a simple `Elemwise` for the subtraction `Op`&#x2013;as expected. There's also an `InplaceDimShuffle` that adds a broadcastable dimension to the second argument, so that both inputs have the same number of dimensions. This `InplaceDimShuffle` is equivalent to `np.expand_dims(m, 0)`, which&#x2013;when subtracted from `x`&#x2013;yields the same value as `z`.

So far, everything is good, because

```python
>>> np.array_equal(aesara.function([], z_at)(), z)
True
```

Now, when we replace the `TensorConstant`s with generic `TensorVariable`s, we get a strange error:

```python
x_v = at.matrix("x")
m_v = at.vector("m")

z_v = x_v - m_v
```

```python
>>> aesara.function([x_v, m_v], z_v)(x, m)
<ipython-input-23-66cc28afa70f> in <module>
----> 1 aesara.function([x_v, m_v], z_v)(x, m)

~/projects/code/python/Aesara/aesara/compile/function/types.py in __call__(self, *args, **kwargs)
    989                     node=self.fn.nodes[self.fn.position_of_error],
    990                     thunk=thunk,
--> 991                     storage_map=getattr(self.fn, "storage_map", None),
    992                 )
    993             else:

~/projects/code/python/Aesara/aesara/link/utils.py in raise_with_op(fgraph, node, thunk, exc_info, storage_map)
    506         # Some exception need extra parameter in inputs. So forget the
    507         # extra long error message in that case.
--> 508     raise exc_value.with_traceback(exc_trace)
    509
    510

~/projects/code/python/Aesara/aesara/compile/function/types.py in __call__(self, *args, **kwargs)
    973             outputs = (
    974                 self.fn()
--> 975                 if output_subset is None
    976                 else self.fn(output_subset=output_subset)
    977             )

ValueError: Input dimension mis-match. (input[0].shape[1] = 1, input[1].shape[1] = 2)
Apply node that caused the error: Elemwise{sub,no_inplace}(x, InplaceDimShuffle{x,0}.0)
Toposort index: 1
Inputs types: [TensorType(float64, matrix), TensorType(float64, row)]
Inputs shapes: [(2, 1), (1, 2)]
Inputs strides: [(8, 8), (16, 8)]
Inputs values: [array([[-1.32720483],
       [ 0.23442016]]), array([[0., 0.]])]
Outputs clients: [['output']]
```

We can emulate this issue using the Python implementation of `Elemwise`, as well:

```python
>>> aesara.function([x_v, m_v], z_v, mode="FAST_COMPILE")(x, m)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
~/projects/code/python/Aesara/aesara/link/vm.py in __call__(self)
    312                 ):
--> 313                     thunk()
    314                     for old_s in old_storage:

~/projects/code/python/Aesara/aesara/graph/op.py in rval(p, i, o, n)
    472             def rval(p=p, i=node_input_storage, o=node_output_storage, n=node):
--> 473                 r = p(n, [x[0] for x in i], o)
    474                 for o in node.outputs:

~/projects/code/python/Aesara/aesara/tensor/elemwise.py in perform(self, node, inputs, output_storage)
    760                 base_exc_str = f"Dimension mismatch; shapes are {', '.join(msg)}"
--> 761                 raise ValueError(base_exc_str)
    762

ValueError: Dimension mismatch; shapes are (2, 1), (*, 2)
```

From this output, we can see that this erroneous error is apparently the result of some bad input validation code in `Elemwise`.

The same is true for the C implementation, although that's a little less apparent from the output. In this case, the C code for this validation step is [here](https://github.com/pymc-devs/aesara/blob/master/aesara/tensor/elemwise_cgen.py#L98).

Comments:

Patch:
diff --git a/aesara/scalar/basic.py b/aesara/scalar/basic.py
index 0bf12a70a..36a05c387 100644
--- a/aesara/scalar/basic.py
+++ b/aesara/scalar/basic.py
@@ -2307,6 +2307,8 @@ class Pow(BinaryScalarOp):
         if (
             node.inputs[0].type == node.outputs[0].type
             and node.inputs[1].type == node.outputs[0].type
+            and None not in node.inputs[0].type.shape
+            and None not in node.inputs[1].type.shape
             and
             # amdlibm 3.0 do not have a float64 version of this SIMD function
             node.inputs[0].dtype == "float32"
diff --git a/aesara/tensor/elemwise.py b/aesara/tensor/elemwise.py
index b678ad16b..ce7d67b5f 100644
--- a/aesara/tensor/elemwise.py
+++ b/aesara/tensor/elemwise.py
@@ -913,7 +913,7 @@ second dimension
         checks = cgen.make_checks(orders, idtypes, sub)
 
         # Check if all inputs (except broadcasted scalar) are fortran.
-        # In that case, create an fortran output ndarray.
+        # In that case, create a fortran output ndarray.
         z = list(zip(inames, inputs))
         alloc_fortran = " && ".join(
             [
@@ -1071,7 +1071,7 @@ second dimension
 
         # If all inputs and outputs are contiguous
         # and the scalar op define optimized code for that case
-        # use it! The scalar_op need to check the broadcast flag himself.
+        # use it! The scalar_op needs to check the type-level shapes itself.
         if (
             all(o.ndim >= 1 for o in node.outputs)
             and
@@ -1088,11 +1088,19 @@ second dimension
                 # compiler to vectorize the code as their won't be as
                 # many ptr and the stride will be hard coded.
                 if all(
-                    [
-                        io.broadcastable == node.outputs[0].broadcastable
-                        or all(io.broadcastable)
-                        for io in node.inputs + node.outputs
-                    ]
+                    # io.type.shape == node.outputs[1].type.shape
+                    # Elemwise does not specify non-broadcastable static/type-levelshape
+                    # information for its outputs yet
+                    node.outputs[0].type.is_super(io.type)
+                    for io in node.inputs + node.outputs
+                ) and (
+                    len(node.inputs) <= 1
+                    # If either one of the inputs has a `None` shape, we cannot
+                    # assume they will have the same size
+                    or all(
+                        len(set(inp_shape)) == 1 and None not in inp_shape
+                        for inp_shape in zip(*(inp.type.shape for inp in node.inputs))
+                    )
                 ):
                     z = onames[0]
                     contig = f"""
@@ -1188,7 +1196,7 @@ second dimension
         return support_code
 
     def c_code_cache_version_apply(self, node):
-        version = [13]  # the version corresponding to the c code in this Op
+        version = [14]  # the version corresponding to the c code in this Op
 
         # now we insert versions for the ops on which we depend...
         scalar_node = Apply(
diff --git a/aesara/tensor/elemwise_cgen.py b/aesara/tensor/elemwise_cgen.py
index cfc5abde8..9e50369fb 100644
--- a/aesara/tensor/elemwise_cgen.py
+++ b/aesara/tensor/elemwise_cgen.py
@@ -66,10 +66,12 @@ def make_checks(loop_orders, dtypes, sub):
             if index != "x":
                 # Initialize the variables associated to the jth loop
                 # jump = stride - adjust
+                # If the variable has size 1 in that dim, we set the stride to zero to
+                # emulate broadcasting
                 jump = f"({var}_stride{index}) - ({adjust})"
                 init += f"""
                 {var}_n{index} = PyArray_DIMS({var})[{index}];
-                {var}_stride{index} = PyArray_STRIDES({var})[{index}] / sizeof({dtype});
+                {var}_stride{index} = ({var}_n{index} == 1)? 0 : PyArray_STRIDES({var})[{index}] / sizeof({dtype});
                 {var}_jump{index}_{j} = {jump};
                 """
                 adjust = f"{var}_n{index}*{var}_stride{index}"
@@ -90,26 +92,85 @@ def make_checks(loop_orders, dtypes, sub):
         # elements of to_compare are pairs ( input_variable_idx, input_variable_dim_idx )
         if len(to_compare) < 2:
             continue
-        j0, x0 = to_compare[0]
-        for (j, x) in to_compare[1:]:
-            check += f"""
-            if (%(lv{j0})s_n{x0} != %(lv{j})s_n{x})
+
+        # Find first dimension size that is != 1
+        jl, xl = to_compare[-1]
+        non1size_dim_check = f"""
+            npy_intp non1size_dim{xl};
+            non1size_dim{xl} = """
+        for (j, x) in to_compare[:-1]:
+            non1size_dim_check += f"(%(lv{j})s_n{x} != 1) ? %(lv{j})s_n{x} : "
+        non1size_dim_check += f"%(lv{jl})s_n{xl};"
+        check += non1size_dim_check
+
+        # Check the nonsize1 dims match
+        # TODO: This is a bit inefficient because we are comparing one dimension against itself
+        check += f"""
+            if (non1size_dim{xl} != 1)
             {{
-                PyErr_Format(PyExc_ValueError, "Input dimension mismatch. (input[%%i].shape[%%i] = %%lld, input[%%i].shape[%%i] = %%lld)",
-                   {j0},
-                   {x0},
-                   (long long int) %(lv{j0})s_n{x0},
-                   {j},
-                   {x},
-                   (long long int) %(lv{j})s_n{x}
-                );
-                %(fail)s
-            }}
+        """
+        for (j, x) in to_compare:
+            check += f"""
+                if ((%(lv{j})s_n{x} != non1size_dim{x}) && (%(lv{j})s_n{x} != 1))
+                {{
+                    PyErr_Format(PyExc_ValueError, "Input dimension mismatch. One other input has shape[%%i] = %%lld, but input[%%i].shape[%%i] = %%lld.",
+                       {x},
+                       (long long int) non1size_dim{x},
+                       {j},
+                       {x},
+                       (long long int) %(lv{j})s_n{x}
+                    );
+                    %(fail)s
+                }}
             """
+        check += """
+            }
+        """
 
     return init % sub + check % sub
 
 
+def compute_broadcast_dimensions(array_name: str, loop_orders, sub) -> str:
+    """Create c_code to compute broadcasted dimensions of multiple arrays, arising from
+    Elemwise operations.
+
+    The code returned by this function populates the array `array_name`, but does not
+    initialize it.
+
+    TODO: We can decide to either specialize C code even further given the input types
+    or make it general, regardless of whether static broadcastable information is given
+    """
+    dims_c_code = ""
+    for i, candidates in enumerate(zip(*loop_orders)):
+        # TODO: Are candidates always either "x" or "i"? If that's the case we can
+        # simplify some logic here (e.g., we don't need to track the `idx`).
+        nonx_candidates = tuple(
+            (idx, c) for idx, c in enumerate(candidates) if c != "x"
+        )
+
+        # All inputs are known to be broadcastable
+        if not nonx_candidates:
+            dims_c_code += f"{array_name}[{i}] = 1;\n"
+            continue
+
+        # There is only one informative source of size
+        if len(nonx_candidates) == 1:
+            idx, candidate = nonx_candidates[0]
+            var = sub[f"lv{int(idx)}"]
+            dims_c_code += f"{array_name}[{i}] = {var}_n{candidate};\n"
+            continue
+
+        # In this case any non-size 1 variable will define the right size
+        dims_c_code += f"{array_name}[{i}] = "
+        for (idx, candidate) in nonx_candidates[:-1]:
+            var = sub[f"lv{int(idx)}"]
+            dims_c_code += f"({var}_n{candidate} != 1)? {var}_n{candidate}: "
+        idx, candidate = nonx_candidates[-1]
+        var = sub[f"lv{idx}"]
+        dims_c_code += f"{var}_n{candidate};\n"
+    return dims_c_code
+
+
 def make_alloc(loop_orders, dtype, sub, fortran="0"):
     """Generate C code to allocate outputs.
 
@@ -125,20 +186,7 @@ def make_alloc(loop_orders, dtype, sub, fortran="0"):
     if type.startswith("AESARA_COMPLEX"):
         type = type.replace("AESARA_COMPLEX", "NPY_COMPLEX")
     nd = len(loop_orders[0])
-    init_dims = ""
-    # For each dimension, the tensors are either all broadcasted, in
-    # which case the output will also be broadcastable (dimension =
-    # 1), or one or more are not broadcasted, in which case the number
-    # of elements of the output in that dimension will be equal to the
-    # number of elements of any of them.
-    for i, candidates in enumerate(zip(*loop_orders)):
-        for j, candidate in enumerate(candidates):
-            if candidate != "x":
-                var = sub[f"lv{int(j)}"]
-                init_dims += f"dims[{i}] = {var}_n{candidate};\n"
-                break
-        else:
-            init_dims += f"dims[{i}] = 1;\n"
+    init_dims = compute_broadcast_dimensions("dims", loop_orders, sub)
 
     # TODO: it would be interesting to allocate the output in such a
     # way that its contiguous dimensions match one of the input's
@@ -310,26 +358,8 @@ def make_reordered_loop(
     """
 
     # Get the (sorted) total number of iterations of each loop
-    # Get totals in the initial order
-    # For each dimension, the tensors are either all broadcasted, in
-    # which case there is only one iteration of the loop, or one or
-    # more are not broadcasted, in which case the number of elements
-    # of any of them will be equal to the number of iterations we have
-    # to do.
-    totals = []
-    for i, candidates in enumerate(zip(*init_loop_orders)):
-        for j, candidate in enumerate(candidates):
-            if candidate != "x":
-                var = sub[f"lv{int(j)}"]
-                total = f"{var}_n{candidate}"
-                break
-        else:
-            total = "1"
-        totals.append(total)
-
-    declare_totals = f"""
-    int init_totals[{nnested}] = {{{", ".join(totals)}}};
-    """
+    declare_totals = f"int init_totals[{nnested}];\n"
+    declare_totals += compute_broadcast_dimensions("init_totals", init_loop_orders, sub)
 
     # Sort totals to match the new order that was computed by sorting
     # the loop vector. One integer variable per loop is declared.


Test patch:
diff --git a/tests/tensor/test_elemwise.py b/tests/tensor/test_elemwise.py
index c2b2d230b..1b298a92a 100644
--- a/tests/tensor/test_elemwise.py
+++ b/tests/tensor/test_elemwise.py
@@ -206,77 +206,117 @@ class TestBroadcast:
         return np.asarray(np.random.random(shp), dtype=aesara.config.floatX)
 
     def with_linker(self, linker, op, type, rand_val):
-        for xsh, ysh in [
-            ((3, 5), (3, 5)),
-            ((3, 5), (1, 5)),
-            ((3, 5), (3, 1)),
-            ((1, 5), (5, 1)),
-            ((1, 1), (1, 1)),
-            ((self.openmp_minsize,), (self.openmp_minsize,)),
-            (
-                (self.openmp_minsize_sqrt, self.openmp_minsize_sqrt),
-                (self.openmp_minsize_sqrt, self.openmp_minsize_sqrt),
-            ),
-            ((2, 3, 4, 5), (2, 3, 4, 5)),
-            ((2, 3, 4, 5), (1, 3, 1, 5)),
-            ((2, 3, 4, 5), (1, 1, 1, 1)),
-            ((), ()),
-        ]:
-            x = type(aesara.config.floatX, [(entry == 1) for entry in xsh])("x")
-            y = type(aesara.config.floatX, [(entry == 1) for entry in ysh])("y")
-            e = op(aes.add)(x, y)
-            f = make_function(copy(linker).accept(FunctionGraph([x, y], [e])))
-            xv = rand_val(xsh)
-            yv = rand_val(ysh)
-            zv = xv + yv
-
-            unittest_tools.assert_allclose(f(xv, yv), zv)
+        for shape_info in ("complete", "only_broadcastable", "none"):
+            for xsh, ysh in [
+                ((3, 5), (3, 5)),
+                ((3, 5), (1, 5)),
+                ((3, 5), (3, 1)),
+                ((1, 5), (5, 1)),
+                ((1, 1), (1, 1)),
+                ((self.openmp_minsize,), (self.openmp_minsize,)),
+                (
+                    (self.openmp_minsize_sqrt, self.openmp_minsize_sqrt),
+                    (self.openmp_minsize_sqrt, self.openmp_minsize_sqrt),
+                ),
+                ((2, 3, 4, 5), (2, 3, 4, 5)),
+                ((2, 3, 4, 5), (1, 3, 1, 5)),
+                ((2, 3, 4, 5), (1, 1, 1, 1)),
+                ((), ()),
+            ]:
+                if shape_info == "complete":
+                    x_type = type(aesara.config.floatX, shape=xsh)
+                    y_type = type(aesara.config.floatX, shape=ysh)
+                elif shape_info == "only_broadcastable":
+                    # This condition is here for backwards compatibility, when the only
+                    # type shape provided by Aesara was broadcastable/non-broadcastable
+                    x_type = type(
+                        aesara.config.floatX,
+                        broadcastable=[(entry == 1) for entry in xsh],
+                    )
+                    y_type = type(
+                        aesara.config.floatX,
+                        broadcastable=[(entry == 1) for entry in ysh],
+                    )
+                else:
+                    x_type = type(aesara.config.floatX, shape=[None for _ in xsh])
+                    y_type = type(aesara.config.floatX, shape=[None for _ in ysh])
 
-            # test Elemwise.infer_shape
-            # the Shape op don't implement c_code!
-            if isinstance(linker, PerformLinker):
-                x = type(aesara.config.floatX, [(entry == 1) for entry in xsh])("x")
-                y = type(aesara.config.floatX, [(entry == 1) for entry in ysh])("y")
+                x = x_type("x")
+                y = y_type("y")
                 e = op(aes.add)(x, y)
-                f = make_function(copy(linker).accept(FunctionGraph([x, y], [e.shape])))
-                assert tuple(f(xv, yv)) == tuple(zv.shape)
+                f = make_function(copy(linker).accept(FunctionGraph([x, y], [e])))
+                xv = rand_val(xsh)
+                yv = rand_val(ysh)
+                zv = xv + yv
 
-    def with_linker_inplace(self, linker, op, type, rand_val):
-        for xsh, ysh in [
-            ((5, 5), (5, 5)),
-            ((5, 5), (1, 5)),
-            ((5, 5), (5, 1)),
-            ((1, 1), (1, 1)),
-            ((2, 3, 4, 5), (2, 3, 4, 5)),
-            ((2, 3, 4, 5), (1, 3, 1, 5)),
-            ((2, 3, 4, 5), (1, 1, 1, 1)),
-            ((), ()),
-        ]:
-            x = type(aesara.config.floatX, [(entry == 1) for entry in xsh])("x")
-            y = type(aesara.config.floatX, [(entry == 1) for entry in ysh])("y")
-            e = op(aes.Add(aes.transfer_type(0)), {0: 0})(x, y)
-            f = make_function(copy(linker).accept(FunctionGraph([x, y], [e])))
-            xv = rand_val(xsh)
-            yv = rand_val(ysh)
-            zv = xv + yv
+                unittest_tools.assert_allclose(f(xv, yv), zv)
 
-            f(xv, yv)
+                # test Elemwise.infer_shape
+                # the Shape op don't implement c_code!
+                if isinstance(linker, PerformLinker):
+                    x = x_type("x")
+                    y = y_type("y")
+                    e = op(aes.add)(x, y)
+                    f = make_function(
+                        copy(linker).accept(FunctionGraph([x, y], [e.shape]))
+                    )
+                    assert tuple(f(xv, yv)) == tuple(zv.shape)
 
-            assert (xv == zv).all()
-            # test Elemwise.infer_shape
-            # the Shape op don't implement c_code!
-            if isinstance(linker, PerformLinker):
-                x = type(aesara.config.floatX, [(entry == 1) for entry in xsh])("x")
-                y = type(aesara.config.floatX, [(entry == 1) for entry in ysh])("y")
+    def with_linker_inplace(self, linker, op, type, rand_val):
+        for shape_info in ("complete", "only_broadcastable", "none"):
+            for xsh, ysh in [
+                ((5, 5), (5, 5)),
+                ((5, 5), (1, 5)),
+                ((5, 5), (5, 1)),
+                ((1, 1), (1, 1)),
+                ((2, 3, 4, 5), (2, 3, 4, 5)),
+                ((2, 3, 4, 5), (1, 3, 1, 5)),
+                ((2, 3, 4, 5), (1, 1, 1, 1)),
+                ((), ()),
+            ]:
+                if shape_info == "complete":
+                    x_type = type(aesara.config.floatX, shape=xsh)
+                    y_type = type(aesara.config.floatX, shape=ysh)
+                elif shape_info == "only_broadcastable":
+                    # This condition is here for backwards compatibility, when the only
+                    # type shape provided by Aesara was broadcastable/non-broadcastable
+                    x_type = type(
+                        aesara.config.floatX,
+                        broadcastable=[(entry == 1) for entry in xsh],
+                    )
+                    y_type = type(
+                        aesara.config.floatX,
+                        broadcastable=[(entry == 1) for entry in ysh],
+                    )
+                else:
+                    x_type = type(aesara.config.floatX, shape=[None for _ in xsh])
+                    y_type = type(aesara.config.floatX, shape=[None for _ in ysh])
+
+                x = x_type("x")
+                y = y_type("y")
                 e = op(aes.Add(aes.transfer_type(0)), {0: 0})(x, y)
-                f = make_function(copy(linker).accept(FunctionGraph([x, y], [e.shape])))
+                f = make_function(copy(linker).accept(FunctionGraph([x, y], [e])))
                 xv = rand_val(xsh)
                 yv = rand_val(ysh)
                 zv = xv + yv
 
                 f(xv, yv)
 
-                assert xv.shape == zv.shape
+                assert (xv == zv).all()
+                # test Elemwise.infer_shape
+                # the Shape op don't implement c_code!
+                if isinstance(linker, PerformLinker):
+                    x = x_type("x")
+                    y = y_type("y")
+                    e = op(aes.Add(aes.transfer_type(0)), {0: 0})(x, y)
+                    f = make_function(
+                        copy(linker).accept(FunctionGraph([x, y], [e.shape]))
+                    )
+                    xv = rand_val(xsh)
+                    yv = rand_val(ysh)
+                    zv = xv + yv
+                    assert xv.shape == zv.shape
+                    assert tuple(f(xv, yv)) == zv.shape
 
     def test_perform(self):
         self.with_linker(PerformLinker(), self.op, self.type, self.rand_val)
@@ -746,10 +786,6 @@ class TestElemwise(unittest_tools.InferShapeTester):
     def test_input_dimensions_match_python(self):
         self.check_input_dimensions_match(Mode(linker="py"))
 
-    @pytest.mark.xfail(
-        reason="Elemwise C implementation does not broadcast parameters",
-        exception=ValueError,
-    )
     @pytest.mark.skipif(
         not aesara.config.cxx, reason="G++ not available, so we need to skip this test."
     )