Cache unique value of TensorConstants and deprecate get_unique_constant_value

Author: ricardoV94

pytensor/scan/rewriting.py

@@ -71,7 +71,7 @@
     get_slice_elements,
     set_subtensor,
 )
-from pytensor.tensor.variable import TensorConstant, get_unique_constant_value
+from pytensor.tensor.variable import TensorConstant
 
 
 list_opt_slice = [
@@ -136,10 +136,7 @@ def remove_constants_and_unused_inputs_scan(fgraph, node):
     all_ins = list(graph_inputs(op_outs))
     for idx in range(op_info.n_seqs):
         node_inp = node.inputs[idx + 1]
-        if (
-            isinstance(node_inp, TensorConstant)
-            and get_unique_constant_value(node_inp) is not None
-        ):
+        if isinstance(node_inp, TensorConstant) and node_inp.unique_value is not None:
             try:
                 # This works if input is a constant that has all entries
                 # equal

pytensor/tensor/basic.py

@@ -19,7 +19,7 @@
 
 import pytensor
 import pytensor.scalar.sharedvar
-from pytensor import compile, config, printing
+from pytensor import config, printing
 from pytensor import scalar as ps
 from pytensor.compile.builders import OpFromGraph
 from pytensor.gradient import DisconnectedType, grad_undefined
@@ -74,7 +74,6 @@
 from pytensor.tensor.variable import (
     TensorConstant,
     TensorVariable,
-    get_unique_constant_value,
 )
 
 
@@ -319,6 +318,8 @@ def get_underlying_scalar_constant_value(
         but I'm not sure where it is.
 
     """
+    from pytensor.compile.ops import DeepCopyOp, OutputGuard
+
     v = orig_v
     while True:
         if v is None:
@@ -336,34 +337,19 @@ def get_underlying_scalar_constant_value(
                 raise NotScalarConstantError()
 
         if isinstance(v, Constant):
-            unique_value = get_unique_constant_value(v)
-            if unique_value is not None:
-                data = unique_value
-            else:
-                data = v.data
-
-            if isinstance(data, np.ndarray):
-                try:
-                    return np.array(data.item(), dtype=v.dtype)
-                except ValueError:
-                    raise NotScalarConstantError()
+            if isinstance(v, TensorConstant) and v.unique_value is not None:
+                return v.unique_value
 
-            from pytensor.sparse.type import SparseTensorType
+            elif isinstance(v, ScalarConstant):
+                return v.data
 
-            if isinstance(v.type, SparseTensorType):
-                raise NotScalarConstantError()
-
-            return data
+            raise NotScalarConstantError()
 
         if not only_process_constants and getattr(v, "owner", None) and max_recur > 0:
             max_recur -= 1
             if isinstance(
                 v.owner.op,
-                Alloc
-                | DimShuffle
-                | Unbroadcast
-                | compile.ops.OutputGuard
-                | compile.DeepCopyOp,
+                Alloc | DimShuffle | Unbroadcast | OutputGuard | DeepCopyOp,
             ):
                 # OutputGuard is only used in debugmode but we
                 # keep it here to avoid problems with old pickles

pytensor/tensor/rewriting/elemwise.py

@@ -41,7 +41,7 @@
     register_specialize,
 )
 from pytensor.tensor.shape import shape_padleft
-from pytensor.tensor.variable import TensorConstant, get_unique_constant_value
+from pytensor.tensor.variable import TensorConstant
 
 
 class InplaceElemwiseOptimizer(GraphRewriter):
@@ -513,7 +513,6 @@ def local_upcast_elemwise_constant_inputs(fgraph, node):
                     new_inputs.append(i)
                 else:
                     try:
-                        # works only for scalars
                         cval_i = get_underlying_scalar_constant_value(
                             i, only_process_constants=True
                         )
@@ -1216,11 +1215,13 @@ def local_inline_composite_constants(fgraph, node):
     inner_replacements = {}
     for outer_inp, inner_inp in zip(node.inputs, composite_op.fgraph.inputs):
         # Complex variables don't have a `c_literal` that can be inlined
-        if "complex" not in outer_inp.type.dtype:
-            unique_value = get_unique_constant_value(outer_inp)
-            if unique_value is not None:
+        if (
+            isinstance(outer_inp, TensorConstant)
+            and "complex" not in outer_inp.type.dtype
+        ):
+            if outer_inp.unique_value is not None:
                 inner_replacements[inner_inp] = ps.constant(
-                    unique_value, dtype=inner_inp.dtype
+                    outer_inp.unique_value, dtype=inner_inp.dtype
                 )
                 continue
         new_outer_inputs.append(outer_inp)

pytensor/tensor/rewriting/math.py

@@ -105,7 +105,6 @@
 from pytensor.tensor.variable import (
     TensorConstant,
     TensorVariable,
-    get_unique_constant_value,
 )
 
 
@@ -137,16 +136,8 @@ def get_constant(v):
         numeric constant. If v is a plain Variable, returns None.
 
     """
-    if isinstance(v, Constant):
-        unique_value = get_unique_constant_value(v)
-        if unique_value is not None:
-            data = unique_value
-        else:
-            data = v.data
-        if data.ndim == 0:
-            return data
-        else:
-            return None
+    if isinstance(v, TensorConstant):
+        return v.unique_value
     elif isinstance(v, Variable):
         return None
     else:

pytensor/tensor/variable.py

@@ -11,7 +11,10 @@
 from pytensor.configdefaults import config
 from pytensor.graph.basic import Constant, OptionalApplyType, Variable
 from pytensor.graph.utils import MetaType
-from pytensor.scalar import ComplexError, IntegerDivisionError
+from pytensor.scalar import (
+    ComplexError,
+    IntegerDivisionError,
+)
 from pytensor.tensor import _get_vector_length
 from pytensor.tensor.exceptions import AdvancedIndexingError
 from pytensor.tensor.type import TensorType
@@ -1042,15 +1045,9 @@ def no_nan(self):
 
 def get_unique_constant_value(x: TensorVariable) -> Number | None:
     """Return the unique value of a tensor, if there is one"""
-    if isinstance(x, Constant):
-        data = x.data
-
-        if isinstance(data, np.ndarray) and data.ndim > 0:
-            flat_data = data.ravel()
-            if flat_data.shape[0]:
-                if (flat_data == flat_data[0]).all():
-                    return flat_data[0]
-
+    warnings.warn("get_unique_constant_value is deprecated.", FutureWarning)
+    if isinstance(x, TensorConstant):
+        return x.unique_value
     return None
 
 
@@ -1077,6 +1074,30 @@ def __init__(self, type: _TensorTypeType, data, name=None):
     def signature(self):
         return TensorConstantSignature((self.type, self.data))
 
+    @property
+    def unique_value(self) -> Number | None:
+        """Return the unique value of a tensor, if there is one"""
+        try:
+            return self._unique_value
+        except AttributeError:
+            data = self.data
+            if np.ndim(data) == 0:
+                unique_value = data
+            else:
+                flat_data = data.ravel()
+                if (flat_data == flat_data[0]).all():
+                    unique_value = flat_data[0]
+                else:
+                    unique_value = None
+
+            if unique_value is not None:
+                # Don't allow the unique value to be changed
+                unique_value.setflags(write=False)
+
+            self._unique_value = unique_value
+
+        return self._unique_value
+
     def equals(self, other):
         # Override Constant.equals to allow to compare with
         # numpy.ndarray, and python type.

tests/tensor/test_basic.py

@@ -3564,11 +3564,12 @@ def test_second(self):
         assert get_underlying_scalar_constant_value(s) == c.data
 
     def test_copy(self):
-        # Make sure we do not return the internal storage of a constant,
+        # Make sure we do not return a writeable internal storage of a constant,
         # so we cannot change the value of a constant by mistake.
         c = constant(3)
         d = extract_constant(c)
-        d += 1
+        with pytest.raises(ValueError, match="output array is read-only"):
+            d += 1
         e = extract_constant(c)
         assert e == 3, (c, d, e)