Address some PR comments

Author: Ch0ronomato

pytensor/link/pytorch/dispatch/blockwise.py

@@ -1,39 +1,35 @@
 import torch
+import torch.compiler
 
 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)
 def funcify_Blockwise(op: Blockwise, node, *args, **kwargs):
     batched_dims = op.batch_ndim(node)
     core_node = op._create_dummy_core_node(node.inputs)
     core_fgraph = FunctionGraph(inputs=core_node.inputs, outputs=core_node.outputs)
-    core_func = pytorch_funcify(core_fgraph)
-    if len(node.outputs) == 1:
-
-        def inner_func(*inputs):
-            return core_func(*inputs)[0]
-    else:
-        inner_func = core_func
+    inner_func = pytorch_funcify(core_fgraph)
 
     for _ in range(batched_dims):
         inner_func = torch.vmap(inner_func)
 
+    @torch.compiler.disable(recursive=False)
     def batcher(*inputs):
         op._check_runtime_broadcast(node, 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,
-        )
+        all_batched_dims = tuple(t.shape[:batched_dims] for t in inputs)
+        batched_shape = torch.broadcast_shapes(*all_batched_dims)
         broadcast_inputs = [
-            torch.broadcast_to(i, s) for i, s in zip(inputs, new_shapes)
+            torch.broadcast_to(i, batched_shape + i.shape[batched_dims:])
+            for i in inputs
         ]
-        return inner_func(*broadcast_inputs)
+        res = inner_func(*broadcast_inputs)
+        if len(node.outputs) == 1:
+            return res[0]
+        else:
+            return res
 
     return batcher

tests/link/pytorch/test_blockwise.py

@@ -3,7 +3,6 @@
 
 import pytensor
 import pytensor.tensor as pt
-from pytensor.tensor.shape import specify_broadcastable
 
 
 torch = pytest.importorskip("torch")
@@ -13,10 +12,10 @@ 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")
+    x = pt.tensor4("x", shape=(5, 1, 2, 3))
+    y = pt.tensor3("y", shape=(3, 3, 2))
 
-    f = pytensor.function([x, y], [x @ y], mode="PYTORCH")
-    [res] = f(_x, _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)