Refactor to get_scalar_constant_value and get_underlying_scalar_constant_value

Author: shreyas3156

doc/library/tensor/basic.rst

@@ -577,7 +577,7 @@ them perfectly, but a `dscalar` otherwise.
     .. method:: round(mode="half_away_from_zero")
        :noindex:
     .. method:: trace()
-    .. method:: get_underlying_scalar_constant()
+    .. method:: get_underlying_scalar_constant_value()
     .. method:: zeros_like(model, dtype=None)
 
        All the above methods are equivalent to NumPy for PyTensor on the current tensor.

pytensor/__init__.py

@@ -153,8 +153,8 @@ def get_underlying_scalar_constant(v):
     if sparse and isinstance(v.type, sparse.SparseTensorType):
         if v.owner is not None and isinstance(v.owner.op, sparse.CSM):
             data = v.owner.inputs[0]
-            return tensor.get_underlying_scalar_constant(data)
-    return tensor.get_underlying_scalar_constant(v)
+            return tensor.get_underlying_scalar_constant_value(data)
+    return tensor.get_underlying_scalar_constant_value(v)
 
 
 # isort: off

pytensor/link/jax/dispatch/tensor_basic.py

@@ -18,7 +18,7 @@
     ScalarFromTensor,
     Split,
     TensorFromScalar,
-    get_underlying_scalar_constant,
+    get_underlying_scalar_constant_value,
 )
 from pytensor.tensor.exceptions import NotScalarConstantError
 
@@ -106,7 +106,7 @@ def join(axis, *tensors):
 def jax_funcify_Split(op: Split, node, **kwargs):
     _, axis, splits = node.inputs
     try:
-        constant_axis = get_underlying_scalar_constant(axis)
+        constant_axis = get_underlying_scalar_constant_value(axis)
     except NotScalarConstantError:
         constant_axis = None
         warnings.warn(
@@ -116,7 +116,7 @@ def jax_funcify_Split(op: Split, node, **kwargs):
     try:
         constant_splits = np.array(
             [
-                get_underlying_scalar_constant(splits[i])
+                get_underlying_scalar_constant_value(splits[i])
                 for i in range(get_vector_length(splits))
             ]
         )

pytensor/scan/basic.py

@@ -12,7 +12,7 @@
 from pytensor.graph.utils import MissingInputError, TestValueError
 from pytensor.scan.op import Scan, ScanInfo
 from pytensor.scan.utils import expand_empty, safe_new, until
-from pytensor.tensor.basic import get_underlying_scalar_constant
+from pytensor.tensor.basic import get_underlying_scalar_constant_value
 from pytensor.tensor.exceptions import NotScalarConstantError
 from pytensor.tensor.math import minimum
 from pytensor.tensor.shape import shape_padleft, unbroadcast
@@ -147,7 +147,7 @@ def isNaN_or_Inf_or_None(x):
         isStr = False
     if not isNaN and not isInf:
         try:
-            val = get_underlying_scalar_constant(x)
+            val = get_underlying_scalar_constant_value(x)
             isInf = np.isinf(val)
             isNaN = np.isnan(val)
         except Exception:
@@ -476,7 +476,7 @@ def wrap_into_list(x):
         n_fixed_steps = int(n_steps)
     else:
         try:
-            n_fixed_steps = at.get_underlying_scalar_constant(n_steps)
+            n_fixed_steps = at.get_underlying_scalar_constant_value(n_steps)
         except NotScalarConstantError:
             n_fixed_steps = None
 

pytensor/scan/rewriting.py

@@ -49,7 +49,11 @@
     safe_new,
     scan_can_remove_outs,
 )
-from pytensor.tensor.basic import Alloc, AllocEmpty, get_underlying_scalar_constant
+from pytensor.tensor.basic import (
+    Alloc,
+    AllocEmpty,
+    get_underlying_scalar_constant_value,
+)
 from pytensor.tensor.elemwise import DimShuffle, Elemwise
 from pytensor.tensor.exceptions import NotScalarConstantError
 from pytensor.tensor.math import Dot, dot, maximum, minimum
@@ -1956,13 +1960,13 @@ def belongs_to_set(self, node, set_nodes):
 
         nsteps = node.inputs[0]
         try:
-            nsteps = int(get_underlying_scalar_constant(nsteps))
+            nsteps = int(get_underlying_scalar_constant_value(nsteps))
         except NotScalarConstantError:
             pass
 
         rep_nsteps = rep.inputs[0]
         try:
-            rep_nsteps = int(get_underlying_scalar_constant(rep_nsteps))
+            rep_nsteps = int(get_underlying_scalar_constant_value(rep_nsteps))
         except NotScalarConstantError:
             pass
 

pytensor/tensor/basic.py

@@ -255,12 +255,14 @@ def _obj_is_wrappable_as_tensor(x):
 )
 
 
-def get_scalar_constant(v, elemwise=True, only_process_constants=False, max_recur=10):
+def get_scalar_constant_value(
+    v, elemwise=True, only_process_constants=False, max_recur=10
+):
     """
     Checks whether 'v' is a scalar (ndim = 0).
 
     If 'v' is a scalar then this function fetches the underlying constant by calling
-    'get_underlying_scalar_constant()'.
+    'get_underlying_scalar_constant_value()'.
 
     If 'v' is not a scalar, it raises a NotScalarConstantError.
 
@@ -269,12 +271,12 @@ def get_scalar_constant(v, elemwise=True, only_process_constants=False, max_recu
         data = v.data
         if data.ndim != 0:
             raise NotScalarConstantError()
-    return get_underlying_scalar_constant(
+    return get_underlying_scalar_constant_value(
         v, elemwise, only_process_constants, max_recur
     )
 
 
-def get_underlying_scalar_constant(
+def get_underlying_scalar_constant_value(
     orig_v, elemwise=True, only_process_constants=False, max_recur=10
 ):
     """Return the constant scalar(0-D) value underlying variable `v`.
@@ -377,7 +379,7 @@ def get_underlying_scalar_constant(
             elif isinstance(v.owner.op, CheckAndRaise):
                 # check if all conditions are constant and true
                 conds = [
-                    get_underlying_scalar_constant(c, max_recur=max_recur)
+                    get_underlying_scalar_constant_value(c, max_recur=max_recur)
                     for c in v.owner.inputs[1:]
                 ]
                 if builtins.all(0 == c.ndim and c != 0 for c in conds):
@@ -391,7 +393,7 @@ def get_underlying_scalar_constant(
                     continue
                 if isinstance(v.owner.op, _scalar_constant_value_elemwise_ops):
                     const = [
-                        get_underlying_scalar_constant(i, max_recur=max_recur)
+                        get_underlying_scalar_constant_value(i, max_recur=max_recur)
                         for i in v.owner.inputs
                     ]
                     ret = [[None]]
@@ -410,7 +412,7 @@ def get_underlying_scalar_constant(
                     v.owner.op.scalar_op, _scalar_constant_value_elemwise_ops
                 ):
                     const = [
-                        get_underlying_scalar_constant(i, max_recur=max_recur)
+                        get_underlying_scalar_constant_value(i, max_recur=max_recur)
                         for i in v.owner.inputs
                     ]
                     ret = [[None]]
@@ -456,7 +458,7 @@ def get_underlying_scalar_constant(
                     ):
                         idx = v.owner.op.idx_list[0]
                         if isinstance(idx, Type):
-                            idx = get_underlying_scalar_constant(
+                            idx = get_underlying_scalar_constant_value(
                                 v.owner.inputs[1], max_recur=max_recur
                             )
                         try:
@@ -490,14 +492,14 @@ def get_underlying_scalar_constant(
                 ):
                     idx = v.owner.op.idx_list[0]
                     if isinstance(idx, Type):
-                        idx = get_underlying_scalar_constant(
+                        idx = get_underlying_scalar_constant_value(
                             v.owner.inputs[1], max_recur=max_recur
                         )
                     # Python 2.4 does not support indexing with numpy.integer
                     # So we cast it.
                     idx = int(idx)
                     ret = v.owner.inputs[0].owner.inputs[idx]
-                    ret = get_underlying_scalar_constant(ret, max_recur=max_recur)
+                    ret = get_underlying_scalar_constant_value(ret, max_recur=max_recur)
                     # MakeVector can cast implicitly its input in some case.
                     return _asarray(ret, dtype=v.type.dtype)
 
@@ -512,7 +514,7 @@ def get_underlying_scalar_constant(
                     idx_list = op.idx_list
                     idx = idx_list[0]
                     if isinstance(idx, Type):
-                        idx = get_underlying_scalar_constant(
+                        idx = get_underlying_scalar_constant_value(
                             owner.inputs[1], max_recur=max_recur
                         )
                     grandparent = leftmost_parent.owner.inputs[0]
@@ -527,7 +529,7 @@ def get_underlying_scalar_constant(
 
                     if not (idx < ndim):
                         msg = (
-                            "get_underlying_scalar_constant detected "
+                            "get_underlying_scalar_constant_value detected "
                             f"deterministic IndexError: x.shape[{int(idx)}] "
                             f"when x.ndim={int(ndim)}."
                         )
@@ -1589,7 +1591,7 @@ def do_constant_folding(self, fgraph, node):
 @_get_vector_length.register(Alloc)
 def _get_vector_length_Alloc(var_inst, var):
     try:
-        return get_underlying_scalar_constant(var.owner.inputs[1])
+        return get_underlying_scalar_constant_value(var.owner.inputs[1])
     except NotScalarConstantError:
         raise ValueError(f"Length of {var} cannot be determined")
 
@@ -1840,17 +1842,17 @@ def perform(self, node, inp, out_):
 
 def extract_constant(x, elemwise=True, only_process_constants=False):
     """
-    This function is basically a call to tensor.get_underlying_scalar_constant.
+    This function is basically a call to tensor.get_underlying_scalar_constant_value.
 
     The main difference is the behaviour in case of failure. While
-    get_underlying_scalar_constant raises an TypeError, this function returns x,
+    get_underlying_scalar_constant_value raises an TypeError, this function returns x,
     as a tensor if possible. If x is a ScalarVariable from a
     scalar_from_tensor, we remove the conversion. If x is just a
     ScalarVariable, we convert it to a tensor with tensor_from_scalar.
 
     """
     try:
-        x = get_underlying_scalar_constant(x, elemwise, only_process_constants)
+        x = get_underlying_scalar_constant_value(x, elemwise, only_process_constants)
     except NotScalarConstantError:
         pass
     if isinstance(x, aes.ScalarVariable) or isinstance(
@@ -2220,7 +2222,7 @@ def make_node(self, axis, *tensors):
 
             if not isinstance(axis, int):
                 try:
-                    axis = int(get_underlying_scalar_constant(axis))
+                    axis = int(get_underlying_scalar_constant_value(axis))
                 except NotScalarConstantError:
                     pass
 
@@ -2469,7 +2471,7 @@ def infer_shape(self, fgraph, node, ishapes):
 def _get_vector_length_Join(op, var):
     axis, *arrays = var.owner.inputs
     try:
-        axis = get_underlying_scalar_constant(axis)
+        axis = get_underlying_scalar_constant_value(axis)
         assert axis == 0 and builtins.all(a.ndim == 1 for a in arrays)
         return builtins.sum(get_vector_length(a) for a in arrays)
     except NotScalarConstantError:
@@ -2881,7 +2883,7 @@ def infer_shape(self, fgraph, node, i_shapes):
 
         def is_constant_value(var, value):
             try:
-                v = get_underlying_scalar_constant(var)
+                v = get_underlying_scalar_constant_value(var)
                 return np.all(v == value)
             except NotScalarConstantError:
                 pass
@@ -4113,8 +4115,8 @@ def take_along_axis(arr, indices, axis=0):
     "cast",
     "scalar_from_tensor",
     "tensor_from_scalar",
-    "get_scalar_constant",
-    "get_underlying_scalar_constant",
+    "get_scalar_constant_value",
+    "get_underlying_scalar_constant_value",
     "constant",
     "as_tensor_variable",
     "as_tensor",

pytensor/tensor/blas.py

@@ -1834,7 +1834,7 @@ def local_gemm_to_ger(fgraph, node):
             xv = x.dimshuffle(0)
             yv = y.dimshuffle(1)
             try:
-                bval = at.get_underlying_scalar_constant(b)
+                bval = at.get_underlying_scalar_constant_value(b)
             except NotScalarConstantError:
                 # b isn't a constant, GEMM is doing useful pre-scaling
                 return

pytensor/tensor/conv/abstract_conv.py

@@ -24,7 +24,10 @@
 from pytensor.graph.basic import Apply, Variable
 from pytensor.graph.op import Op
 from pytensor.raise_op import Assert
-from pytensor.tensor.basic import as_tensor_variable, get_underlying_scalar_constant
+from pytensor.tensor.basic import (
+    as_tensor_variable,
+    get_underlying_scalar_constant_value,
+)
 from pytensor.tensor.exceptions import NotScalarConstantError
 from pytensor.tensor.var import TensorConstant, TensorVariable
 
@@ -495,8 +498,8 @@ def check_dim(given, computed):
         if given is None or computed is None:
             return True
         try:
-            given = get_underlying_scalar_constant(given)
-            computed = get_underlying_scalar_constant(computed)
+            given = get_underlying_scalar_constant_value(given)
+            computed = get_underlying_scalar_constant_value(computed)
             return int(given) == int(computed)
         except NotScalarConstantError:
             # no answer possible, accept for now
@@ -532,7 +535,7 @@ def assert_conv_shape(shape):
     out_shape = []
     for i, n in enumerate(shape):
         try:
-            const_n = get_underlying_scalar_constant(n)
+            const_n = get_underlying_scalar_constant_value(n)
             if i < 2:
                 if const_n < 0:
                     raise ValueError(
@@ -2200,7 +2203,9 @@ def __init__(
             if imshp_i is not None:
                 # Components of imshp should be constant or ints
                 try:
-                    get_underlying_scalar_constant(imshp_i, only_process_constants=True)
+                    get_underlying_scalar_constant_value(
+                        imshp_i, only_process_constants=True
+                    )
                 except NotScalarConstantError:
                     raise ValueError(
                         "imshp should be None or a tuple of constant int values"
@@ -2213,7 +2218,9 @@ def __init__(
             if kshp_i is not None:
                 # Components of kshp should be constant or ints
                 try:
-                    get_underlying_scalar_constant(kshp_i, only_process_constants=True)
+                    get_underlying_scalar_constant_value(
+                        kshp_i, only_process_constants=True
+                    )
                 except NotScalarConstantError:
                     raise ValueError(
                         "kshp should be None or a tuple of constant int values"

pytensor/tensor/elemwise.py

@@ -759,7 +759,7 @@ def perform(self, node, inputs, output_storage):
                 ufunc = self.ufunc
             elif not hasattr(node.tag, "ufunc"):
                 # It happen that make_thunk isn't called, like in
-                # get_underlying_scalar_constant
+                # get_underlying_scalar_constant_value
                 self.prepare_node(node, None, None, "py")
                 # prepare_node will add ufunc to self or the tag
                 # depending if we can reuse it or not. So we need to

pytensor/tensor/exceptions.py

@@ -4,7 +4,7 @@ class ShapeError(Exception):
 
 class NotScalarConstantError(Exception):
     """
-    Raised by get_underlying_scalar_constant if called on something that is
+    Raised by get_underlying_scalar_constant_value if called on something that is
     not a scalar constant.
     """
 

pytensor/tensor/extra_ops.py

@@ -671,7 +671,7 @@ def make_node(self, x, repeats):
             out_shape = [None]
         else:
             try:
-                const_reps = at.get_underlying_scalar_constant(repeats)
+                const_reps = at.get_underlying_scalar_constant_value(repeats)
             except NotScalarConstantError:
                 const_reps = None
             if const_reps == 1:

pytensor/tensor/random/op.py

@@ -12,7 +12,7 @@
 from pytensor.tensor.basic import (
     as_tensor_variable,
     constant,
-    get_underlying_scalar_constant,
+    get_underlying_scalar_constant_value,
     get_vector_length,
     infer_static_shape,
 )
@@ -277,7 +277,7 @@ def infer_shape(self, fgraph, node, input_shapes):
         try:
             size_len = get_vector_length(size)
         except ValueError:
-            size_len = get_underlying_scalar_constant(size_shape[0])
+            size_len = get_underlying_scalar_constant_value(size_shape[0])
 
         size = tuple(size[n] for n in range(size_len))