Ensure inputs are shaped

Author: Ch0ronomato

pytensor/link/pytorch/dispatch/blockwise.py

@@ -3,6 +3,7 @@
 from pytensor.graph import FunctionGraph
 from pytensor.link.pytorch.dispatch import pytorch_funcify
 from pytensor.tensor.blockwise import Blockwise
+from pytensor.tensor.random.utils import params_broadcast_shapes
 
 
 @pytorch_funcify.register(Blockwise)
@@ -23,6 +24,16 @@ def inner_func(*inputs):
 
     def batcher(*inputs):
         op._check_runtime_broadcast(node, inputs)
-        return inner_func(*inputs)
+        # broadcast on batched_dims
+        all_batched_dims = tuple(tuple(t.shape) for t in inputs)
+        new_shapes = params_broadcast_shapes(
+            all_batched_dims,
+            ndims_params=[batched_dims] * len(inputs),
+            use_pytensor=False,
+        )
+        broadcast_inputs = [
+            torch.broadcast_to(i, s) for i, s in zip(inputs, new_shapes)
+        ]
+        return inner_func(*broadcast_inputs)
 
     return batcher

tests/link/pytorch/test_blockwise.py

@@ -1,9 +1,13 @@
+import numpy as np
 import pytest
 
+import pytensor
+import pytensor.tensor as pt
 from pytensor.graph.replace import vectorize_node
 from pytensor.tensor import tensor
 from pytensor.tensor.blockwise import Blockwise
 from pytensor.tensor.nlinalg import MatrixInverse
+from pytensor.tensor.shape import specify_broadcastable
 
 
 torch = pytest.importorskip("torch")
@@ -30,3 +34,16 @@ def test_vectorize_blockwise():
         new_vect_node.op.core_op, MatrixInverse
     )
     assert new_vect_node.inputs[0] is tns4
+
+
+def test_blockwise_broadcast():
+    _x = np.random.rand(5, 1, 2, 3)
+    _y = np.random.rand(3, 3, 2)
+
+    x = specify_broadcastable(pt.tensor4("x"), 1)
+    y = pt.tensor3("y")
+
+    f = pytensor.function([x, y], [x @ y], mode="PYTORCH")
+    [res] = f(_x, _y)
+    assert tuple(res.shape) == (5, 3, 2, 2)
+    np.testing.assert_allclose(res, _x @ _y)