Make blockwise perform method node dependent

Author: ricardoV94

pytensor/tensor/blockwise.py

@@ -1,5 +1,4 @@
 from collections.abc import Sequence
-from copy import copy
 from typing import Any, cast
 
 import numpy as np
@@ -79,7 +78,6 @@ def __init__(
         self.name = name
         self.inputs_sig, self.outputs_sig = _parse_gufunc_signature(signature)
         self.gufunc_spec = gufunc_spec
-        self._gufunc = None
         if destroy_map is not None:
             self.destroy_map = destroy_map
         if self.destroy_map != core_op.destroy_map:
@@ -91,11 +89,6 @@ def __init__(
 
         super().__init__(**kwargs)
 
-    def __getstate__(self):
-        d = copy(self.__dict__)
-        d["_gufunc"] = None
-        return d
-
     def _create_dummy_core_node(self, inputs: Sequence[TensorVariable]) -> Apply:
         core_input_types = []
         for i, (inp, sig) in enumerate(zip(inputs, self.inputs_sig)):
@@ -320,8 +313,7 @@ def core_func(*inner_inputs):
             else:
                 return tuple(r[0] for r in inner_outputs)
 
-        self._gufunc = np.vectorize(core_func, signature=self.signature)
-        return self._gufunc
+        node.tag.gufunc = np.vectorize(core_func, signature=self.signature)
 
     def _check_runtime_broadcast(self, node, inputs):
         batch_ndim = self.batch_ndim(node)
@@ -340,10 +332,12 @@ def _check_runtime_broadcast(self, node, inputs):
                 )
 
     def perform(self, node, inputs, output_storage):
-        gufunc = self._gufunc
+        gufunc = getattr(node.tag, "gufunc", None)
 
         if gufunc is None:
-            gufunc = self._create_gufunc(node)
+            # Cache it once per node
+            self._create_gufunc(node)
+            gufunc = node.tag.gufunc
 
         self._check_runtime_broadcast(node, inputs)
 

tests/tensor/test_blockwise.py

@@ -28,6 +28,41 @@
 from pytensor.tensor.utils import _parse_gufunc_signature
 
 
+def test_perform_method_per_node():
+    """Confirm that Blockwise uses one perform method per node.
+
+    This is important if the perform method requires node information (such as dtypes)
+    """
+
+    class NodeDependentPerformOp(Op):
+        def make_node(self, x):
+            return Apply(self, [x], [x.type()])
+
+        def perform(self, node, inputs, outputs):
+            [x] = inputs
+            if node.inputs[0].type.dtype.startswith("float"):
+                y = x + 1
+            else:
+                y = x - 1
+            outputs[0][0] = y
+
+    blockwise_op = Blockwise(core_op=NodeDependentPerformOp(), signature="()->()")
+    x = tensor("x", shape=(3,), dtype="float32")
+    y = tensor("y", shape=(3,), dtype="int32")
+
+    out_x = blockwise_op(x)
+    out_y = blockwise_op(y)
+    fn = pytensor.function([x, y], [out_x, out_y])
+    [op1, op2] = [node.op for node in fn.maker.fgraph.apply_nodes]
+    # Confirm both nodes have the same Op
+    assert op1 is blockwise_op
+    assert op1 is op2
+
+    res_out_x, res_out_y = fn(np.zeros(3, dtype="float32"), np.zeros(3, dtype="int32"))
+    np.testing.assert_array_equal(res_out_x, np.ones(3, dtype="float32"))
+    np.testing.assert_array_equal(res_out_y, -np.ones(3, dtype="int32"))
+
+
 def test_vectorize_blockwise():
     mat = tensor(shape=(None, None))
     tns = tensor(shape=(None, None, None))