Handle invalid BroadcastTo shape in C code

Author: ricardoV94

pytensor/tensor/extra_ops.py

@@ -1643,6 +1643,9 @@ def make_node(self, a, *shape):
 
         shape, static_shape = at.infer_static_shape(shape)
 
+        if len(shape) < a.ndim:
+            raise ValueError("Broadcast shape cannot be shorter than input shape")
+
         out = TensorType(dtype=a.type.dtype, shape=static_shape)()
 
         # Attempt to prevent in-place operations on this view-based output
@@ -1715,6 +1718,23 @@ def c_code(self, node, name, inputs, outputs, sub):
                 1, ops, flags, NPY_CORDER, NPY_NO_CASTING, op_flags, op_dtypes, oa_ndim, op_axes, itershape, buffersize
             );
 
+
+            int offset = %(ndims)s - PyArray_NDIM(%(x)s);
+
+            for(int i = 0; i < PyArray_NDIM(%(x)s); i++)
+            {
+                if ((PyArray_DIMS(%(x)s)[i] != 1) && (PyArray_DIMS(%(x)s)[i] != itershape[i + offset]))
+                {
+                    PyErr_Format(PyExc_ValueError,
+                                 "Shape mismatch in broadcast_to: target shape[%%i] = %%lld is incompatible with input shape = %%lld.",
+                                 i,
+                                 (long long int) itershape[i + offset],
+                                 (long long int) PyArray_DIMS(%(x)s)[i]
+                    );
+                    %(fail)s
+                }
+            }
+
             %(out)s = NpyIter_GetIterView(iter, 0);
 
             if(%(out)s == NULL){
@@ -1733,7 +1753,7 @@ def c_code(self, node, name, inputs, outputs, sub):
         return src
 
     def c_code_cache_version(self):
-        return (1,)
+        return (2,)
 
 
 broadcast_to_ = BroadcastTo()

tests/tensor/test_extra_ops.py

@@ -1253,30 +1253,39 @@ def test_avoid_useless_subtensors(self):
     @pytest.mark.parametrize("linker", ["cvm", "py"])
     def test_perform(self, linker):
 
-        a = pytensor.shared(5)
+        a = pytensor.shared(np.full((3, 1, 1), 5))
+        s_0 = iscalar("s_0")
         s_1 = iscalar("s_1")
-        shape = (s_1, 1)
+        shape = (s_0, s_1, 1)
 
         bcast_res = broadcast_to(a, shape)
-        assert bcast_res.broadcastable == (False, True)
+        assert bcast_res.broadcastable == (False, False, True)
 
         bcast_fn = pytensor.function(
-            [s_1], bcast_res, mode=Mode(optimizer=None, linker=linker)
+            [s_0, s_1], bcast_res, mode=Mode(optimizer=None, linker=linker)
         )
         bcast_fn.vm.allow_gc = False
 
-        bcast_at = bcast_fn(4)
-        bcast_np = np.broadcast_to(5, (4, 1))
+        bcast_at = bcast_fn(3, 4)
+        bcast_np = np.broadcast_to(5, (3, 4, 1))
 
         assert np.array_equal(bcast_at, bcast_np)
 
-        bcast_var = bcast_fn.maker.fgraph.outputs[0].owner.inputs[0]
-        bcast_in = bcast_fn.vm.storage_map[a]
-        bcast_out = bcast_fn.vm.storage_map[bcast_var]
+        with pytest.raises(ValueError):
+            bcast_fn(5, 4)
 
         if linker != "py":
+            bcast_var = bcast_fn.maker.fgraph.outputs[0].owner.inputs[0]
+            bcast_in = bcast_fn.vm.storage_map[a]
+            bcast_out = bcast_fn.vm.storage_map[bcast_var]
             assert np.shares_memory(bcast_out[0], bcast_in[0])
 
+    def test_make_node_error_handling(self):
+        with pytest.raises(
+            ValueError, match="Broadcast shape cannot be shorter than input shape"
+        ):
+            broadcast_to(at.zeros((3, 4)), (5,))
+
     @pytest.mark.skipif(
         not config.cxx, reason="G++ not available, so we need to skip this test."
     )