Numba Alloc: Patch so it works inside a Blockwise

Author: ricardoV94

pytensor/link/numba/dispatch/tensor_basic.py

@@ -59,7 +59,7 @@ def numba_funcify_Alloc(op, node, **kwargs):
     shape_var_item_names = [f"{name}_item" for name in shape_var_names]
     shapes_to_items_src = indent(
         "\n".join(
-            f"{item_name} = {shape_name}.item()"
+            f"{item_name} = int({shape_name})"
             for item_name, shape_name in zip(
                 shape_var_item_names, shape_var_names, strict=True
             )
@@ -74,24 +74,26 @@ def numba_funcify_Alloc(op, node, **kwargs):
                 f'if val.shape[{-i - 1}] == 1 and scalar_shape[{-i - 1}] != 1: raise ValueError("{Alloc._runtime_broadcast_error_msg}")'
             )
     check_runtime_broadcast_src = indent("\n".join(check_runtime_broadcast), " " * 4)
-
+    dtype = node.inputs[0].type.dtype
     alloc_def_src = f"""
 def alloc(val, {", ".join(shape_var_names)}):
 {shapes_to_items_src}
     scalar_shape = {create_tuple_string(shape_var_item_names)}
 {check_runtime_broadcast_src}
-    res = np.empty(scalar_shape, dtype=val.dtype)
+    res = np.empty(scalar_shape, dtype=np.{dtype})
     res[...] = val
     return res
     """
     alloc_fn = compile_numba_function_src(
         alloc_def_src,
         "alloc",
         globals() | {"np": np},
+        write_to_disk=True,
     )
 
+    cache_version = -1
     cache_key = sha256(
-        str((type(op), node.inputs[0].type.broadcastable)).encode()
+        str((type(op), node.inputs[0].type.broadcastable, cache_version)).encode()
     ).hexdigest()
     return numba_basic.numba_njit(alloc_fn), cache_key
 

tests/link/numba/test_blockwise.py

@@ -2,8 +2,8 @@
 import pytest
 
 from pytensor import function
-from pytensor.tensor import tensor, tensor3
-from pytensor.tensor.basic import ARange
+from pytensor.tensor import lvector, tensor, tensor3
+from pytensor.tensor.basic import Alloc, ARange, constant
 from pytensor.tensor.blockwise import Blockwise, BlockwiseWithCoreShape
 from pytensor.tensor.nlinalg import SVD, Det
 from pytensor.tensor.slinalg import Cholesky, cholesky
@@ -70,3 +70,13 @@ def test_repeated_args():
     final_node = fn.maker.fgraph.outputs[0].owner
     assert isinstance(final_node.op, BlockwiseWithCoreShape)
     assert final_node.inputs[0] is final_node.inputs[1]
+
+
+def test_blockwise_alloc():
+    val = lvector("val")
+    out = Blockwise(Alloc(), signature="(),(),()->(2,3)")(
+        val, constant(2, dtype="int64"), constant(3, dtype="int64")
+    )
+    assert out.type.ndim == 3
+
+    compare_numba_and_py([val], [out], [np.arange(5)], eval_obj_mode=False)