Flag Ops whose output types depend on input values

These nodes must always be rebuilt in non-strict mode

Author: ricardoV94

pytensor/graph/basic.py

@@ -266,8 +266,15 @@ def clone_with_new_inputs(
         assert isinstance(inputs, (list, tuple))
         remake_node = False
         new_inputs: List["Variable"] = list(inputs)
+
+        # Some Ops like Alloc require the node to always be rebuilt
+        # as the output type depends on the input values and not just their types
+        output_type_depends_on_input_value = getattr(
+            self.op, "_output_type_depends_on_input_value", False
+        )
+
         for i, (curr, new) in enumerate(zip(self.inputs, new_inputs)):
-            if curr.type != new.type:
+            if (curr.type != new.type) or output_type_depends_on_input_value:
                 if strict:
                     new_i = curr.type.filter_variable(new)
                     new_inputs[i] = new_i

pytensor/tensor/basic.py

@@ -1418,6 +1418,8 @@ class Alloc(COp):
     """
 
     _f16_ok = True
+    _output_type_depends_on_input_value = True
+
     __props__ = ()
 
     def make_node(self, value, *shape):
@@ -3817,6 +3819,8 @@ def perform(self, node, inputs, outputs):
 class AllocEmpty(COp):
     """Implement Alloc on the cpu, but without initializing memory."""
 
+    _output_type_depends_on_input_value = True
+
     __props__ = ("dtype",)
     params_type = ParamsType(typecode=int32)
 

pytensor/tensor/extra_ops.py

@@ -1630,6 +1630,8 @@ def broadcast_shape_iter(
 class BroadcastTo(COp):
     """An `Op` for `numpy.broadcast_to`."""
 
+    _output_type_depends_on_input_value = True
+
     __props__ = ()
 
     view_map = {0: [0]}

pytensor/tensor/random/op.py

@@ -91,6 +91,8 @@ class RandomVariable(Op):
 
     """
 
+    _output_type_depends_on_input_value = True
+
     __props__ = ("name", "ndim_supp", "ndims_params", "dtype", "inplace")
     default_output = 1
 

tests/tensor/random/test_basic.py

@@ -14,6 +14,7 @@
 from pytensor.graph.basic import Constant, Variable, graph_inputs
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.op import get_test_value
+from pytensor.graph.replace import clone_replace
 from pytensor.graph.rewriting.db import RewriteDatabaseQuery
 from pytensor.tensor.random.basic import (
     bernoulli,
@@ -57,7 +58,7 @@
     weibull,
 )
 from pytensor.tensor.rewriting.shape import ShapeFeature
-from pytensor.tensor.type import iscalar, scalar, tensor
+from pytensor.tensor.type import iscalar, scalar, vector, tensor
 from tests.unittest_tools import create_pytensor_param
 
 
@@ -1422,3 +1423,17 @@ def test_pickle():
     a_unpkl = pickle.loads(a_pkl)
 
     assert a_unpkl.owner.op._props() == sample_a.owner.op._props()
+
+
+def test_rebuild():
+    x = vector(shape=(50,))
+    x_test = np.zeros((50,))
+    y = normal(size=x.shape)
+    assert y.shape.eval({x: x_test}) == (50,)
+    assert y.eval({x: x_test}).shape == (50,)
+
+    x_new = vector(shape=(100,))
+    x_new_test = np.zeros((100,))
+    y_new = clone_replace(y, {x: x_new}, rebuild_strict=True)
+    assert y_new.shape.eval({x_new: x_new_test}) == (100,)
+    assert y_new.eval({x_new: x_new_test}).shape == (100,)

tests/tensor/test_basic.py

@@ -16,6 +16,7 @@
 from pytensor.gradient import grad, hessian
 from pytensor.graph.basic import Apply
 from pytensor.graph.op import Op
+from pytensor.graph.replace import clone_replace
 from pytensor.misc.safe_asarray import _asarray
 from pytensor.raise_op import Assert
 from pytensor.scalar import autocast_float, autocast_float_as
@@ -818,6 +819,20 @@ def test_full(self):
         res = pytensor.function([], full_at, mode=self.mode)()
         assert np.array_equal(res, np.full((2, 3), 3, dtype="int64"))
 
+    @pytest.mark.parametrize("func", (at.zeros, at.empty))
+    def test_rebuild(self, func):
+        x = vector(shape=(50,))
+        x_test = np.zeros((50,))
+        y = func(x.shape)
+        assert y.shape.eval({x: x_test}) == (50,)
+        assert y.eval({x: x_test}).shape == (50,)
+
+        x_new = vector(shape=(100,))
+        x_new_test = np.zeros((100,))
+        y_new = clone_replace(y, {x: x_new}, rebuild_strict=False)
+        assert y_new.shape.eval({x_new: x_new_test}) == (100,)
+        assert y_new.eval({x_new: x_new_test}).shape == (100,)
+
 
 def test_infer_shape():
     with pytest.raises(TypeError, match="^Shapes must be scalar integers.*"):

tests/tensor/test_extra_ops.py

@@ -10,6 +10,7 @@
 from pytensor.configdefaults import config
 from pytensor.graph.basic import Constant, applys_between
 from pytensor.graph.rewriting.db import RewriteDatabaseQuery
+from pytensor.graph.replace import clone_replace
 from pytensor.raise_op import Assert
 from pytensor.tensor.elemwise import DimShuffle
 from pytensor.tensor.extra_ops import (
@@ -1393,6 +1394,20 @@ def test_inplace(self):
 
         assert advincsub_node.op.inplace is False
 
+    def test_rebuild(self):
+        x = vector(shape=(50,))
+        x_test = np.zeros((50,))
+        i = 0
+        y = broadcast_to(i, x.shape)
+        assert y.shape.eval({x: x_test}) == (50,)
+        assert y.eval({x: x_test}).shape == (50,)
+
+        x_new = vector(shape=(100,))
+        x_new_test = np.zeros((100,))
+        y_new = clone_replace(y, {x: x_new}, rebuild_strict=False)
+        assert y_new.shape.eval({x_new: x_new_test}) == (100,)
+        assert y_new.eval({x_new: x_new_test}).shape == (100,)
+
 
 def test_broadcast_arrays():
     x, y = at.dvector(), at.dmatrix()