pytensor 2.38.1__tar.gz → 2.38.2__tar.gz

{pytensor-2.38.1/pytensor.egg-info → pytensor-2.38.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 2.38.1
+Version: 2.38.2
 Summary: Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs.
 Author-email: pymc-devs <pymc.devs@gmail.com>
 License-Expression: BSD-3-Clause

{pytensor-2.38.1 → pytensor-2.38.2}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2026-03-01T12:50:03+0000",
+ "date": "2026-03-06T13:58:13+0100",
  "dirty": false,
  "error": null,
- "full-revisionid": "ecd1e07304754e49bcedeeb5ba2ca0e8e9bf4959",
- "version": "2.38.1"
+ "full-revisionid": "7fa753b583fd16f46754b356c21055c9559080ef",
+ "version": "2.38.2"
 }
 '''  # END VERSION_JSON
 

{pytensor-2.38.1 → pytensor-2.38.2}/pytensor/tensor/basic.py

@@ -29,13 +29,13 @@ from pytensor.graph.fg import FunctionGraph, Output
 from pytensor.graph.op import Op
 from pytensor.graph.replace import _vectorize_node
 from pytensor.graph.rewriting.db import EquilibriumDB
-from pytensor.graph.type import HasShape, Type
+from pytensor.graph.type import HasDataType, HasShape, Type
 from pytensor.link.c.op import COp
 from pytensor.link.c.params_type import ParamsType
 from pytensor.printing import Printer, min_informative_str, pprint, set_precedence
 from pytensor.raise_op import CheckAndRaise
 from pytensor.scalar import int32
-from pytensor.scalar.basic import ScalarConstant, ScalarType, ScalarVariable
+from pytensor.scalar.basic import ScalarConstant, ScalarVariable
 from pytensor.tensor import (
     _as_tensor_variable,
     _get_vector_length,
@@ -292,13 +292,8 @@ def _get_underlying_scalar_constant_value(
     max_recur : int
         The maximum number of recursion.
 
-    Notes
-    -----
-        There may be another function similar to this one in the code,
-        but I'm not sure where it is.
-
     """
-    from pytensor.compile.ops import DeepCopyOp, OutputGuard
+    from pytensor.compile.ops import DeepCopyOp, OutputGuard, TypeCastingOp
     from pytensor.sparse import CSM
     from pytensor.tensor.subtensor import Subtensor
 
@@ -319,13 +314,20 @@ def _get_underlying_scalar_constant_value(
                 raise NotScalarConstantError()
 
         if isinstance(v, Constant):
-            if isinstance(v.type, TensorType) and v.unique_value is not None:
-                return v.unique_value
+            v_type = v.type
+            if isinstance(v_type, HasShape) and isinstance(v_type, HasDataType):
+                if v_type.ndim == 0:
+                    return np.array(v.data, dtype=v.type.dtype)
 
-            elif isinstance(v.type, ScalarType):
-                return v.data
+                elif (not any(s is None for s in v_type.shape)) and (
+                    np.prod(v_type.shape) == 1
+                ):
+                    return np.array(v.data, dtype=v_type.dtype).squeeze()
 
-            elif isinstance(v.type, NoneTypeT):
+                elif isinstance(v_type, TensorType) and v.unique_value is not None:
+                    return np.array(v.unique_value)
+
+            elif isinstance(v_type, NoneTypeT):
                 return None
 
             raise NotScalarConstantError()
@@ -333,9 +335,9 @@ def _get_underlying_scalar_constant_value(
         if not only_process_constants and getattr(v, "owner", None) and max_recur > 0:
             op = v.owner.op
             max_recur -= 1
-            if isinstance(op, Alloc | DimShuffle | OutputGuard | DeepCopyOp):
-                # OutputGuard is only used in debugmode but we
-                # keep it here to avoid problems with old pickles
+            if isinstance(
+                op, Alloc | DimShuffle | TypeCastingOp | DeepCopyOp | OutputGuard
+            ):
                 v = v.owner.inputs[0]
                 continue
             elif isinstance(op, Shape_i):
@@ -343,7 +345,6 @@ def _get_underlying_scalar_constant_value(
                 inp = v.owner.inputs[0]
                 if isinstance(inp, Constant):
                     return np.asarray(np.shape(inp.data)[i])
-                # The shape of a broadcastable dimension is 1
                 if isinstance(inp.type, HasShape) and inp.type.shape[i] is not None:
                     return np.asarray(inp.type.shape[i])
 
@@ -600,7 +601,10 @@ def get_scalar_constant_value(
     If 'v' is not a scalar, it raises a NotScalarConstantError.
 
     """
-    if isinstance(v, TensorVariable | np.ndarray):
+    if isinstance(v, Variable) and isinstance(v.type, HasShape):
+        if v.type.ndim != 0:
+            raise NotScalarConstantError("Input ndim != 0")
+    elif isinstance(v, np.ndarray):
         if v.ndim != 0:
             raise NotScalarConstantError("Input ndim != 0")
     return get_underlying_scalar_constant_value(

{pytensor-2.38.1 → pytensor-2.38.2}/pytensor/tensor/subtensor.py

@@ -805,7 +805,7 @@ def get_constant_idx(
     >>> b.owner.op.idx_list
     (ScalarType(int64), slice(ScalarType(int64), ScalarType(int64), None))
     >>> get_constant_idx(b.owner.op.idx_list, b.owner.inputs, allow_partial=True)
-    [v, slice(np.int64(1), np.int64(3), None)]
+    [v, slice(1, 3, None)]
     >>> get_constant_idx(b.owner.op.idx_list, b.owner.inputs)
     Traceback (most recent call last):
     pytensor.tensor.exceptions.NotScalarConstantError
@@ -825,7 +825,7 @@ def get_constant_idx(
                     val,
                     only_process_constants=only_process_constants,
                     elemwise=elemwise,
-                )
+                ).item()
             except NotScalarConstantError:
                 if allow_partial:
                     return val

{pytensor-2.38.1 → pytensor-2.38.2}/pytensor/xtensor/rewriting/shape.py

@@ -119,7 +119,6 @@ def lower_expand_dims(fgraph, node):
 
     # Convert inputs to tensors
     x_tensor = tensor_from_xtensor(x)
-    size_tensor = tensor_from_xtensor(size)
 
     # Get the new dimension name and position
     new_axis = 0  # Always insert at front
@@ -130,7 +129,7 @@ def lower_expand_dims(fgraph, node):
         result_tensor = expand_dims(x_tensor, new_axis)
     else:
         # Otherwise broadcast to the requested size
-        result_tensor = broadcast_to(x_tensor, (size_tensor, *x_tensor.shape))
+        result_tensor = broadcast_to(x_tensor, (size, *x_tensor.shape))
 
     # Preserve static shape information
     result_tensor = specify_shape(result_tensor, out.type.shape)

{pytensor-2.38.1 → pytensor-2.38.2}/pytensor/xtensor/shape.py

@@ -123,7 +123,10 @@ class UnStack(XOp):
             raise ValueError(
                 f"Number of unstacked lengths {len(unstacked_length)} must match number of unstacked dims {len(self.unstacked_dims)}"
             )
-        unstacked_lengths = [as_tensor(length, ndim=0) for length in unstacked_length]
+        unstacked_lengths = [
+            as_tensor(length, allow_xtensor_conversion=True)
+            for length in unstacked_length
+        ]
         if not all(length.dtype in discrete_dtypes for length in unstacked_lengths):
             raise TypeError("Unstacked lengths must be discrete dtypes.")
 
@@ -441,7 +444,7 @@ class ExpandDims(XOp):
         if self.dim in x.type.dims:
             raise ValueError(f"Dimension {self.dim} already exists in {x.type.dims}")
 
-        size = as_xtensor(size, dims=())
+        size = as_tensor(size, allow_xtensor_conversion=True)
         if not (size.dtype in integer_dtypes and size.ndim == 0):
             raise ValueError(f"size should be an integer scalar, got {size.type}")
         try:

{pytensor-2.38.1 → pytensor-2.38.2}/pytensor/xtensor/vectorization.py

@@ -16,6 +16,7 @@ from pytensor.graph.type import HasShape
 from pytensor.scalar import discrete_dtypes
 from pytensor.tensor import (
     TensorVariable,
+    as_tensor,
     broadcast_shape,
     broadcast_to,
     tensor,
@@ -198,6 +199,8 @@ class XRV(XOp, RNGConsumerOp):
         if len(set(extra_dims)) != len(extra_dims):
             raise ValueError("size_dims must be unique")
         self.extra_dims = tuple(extra_dims)
+        if print_name := getattr(core_op, "_print_name", None):
+            self._print_name = print_name
 
     def __str__(self):
         if self.name is not None:
@@ -208,6 +211,15 @@ class XRV(XOp, RNGConsumerOp):
             attrs = f"(core_op={self.core_op}, core_dims={self.core_dims}, extra_dims={self.extra_dims})"
         return f"{name}({attrs})"
 
+    def rng_param(self, node):
+        return node.inputs[0]
+
+    def size_params(self, node):
+        return node.inputs[1 : 1 + len(self.extra_dims)]
+
+    def dist_params(self, node):
+        return node.inputs[1 + len(self.extra_dims) :]
+
     def update(self, node):
         # RNG input and update are the first input and output respectively
         return {node.inputs[0]: node.outputs[0]}
@@ -221,7 +233,7 @@ class XRV(XOp, RNGConsumerOp):
             )
 
         extra_dim_lengths = [
-            as_xtensor(dim_length).values
+            as_tensor(dim_length, allow_xtensor_conversion=True)
             for dim_length in extra_dim_lengths_and_params[: len(self.extra_dims)]
         ]
         if not all(

{pytensor-2.38.1 → pytensor-2.38.2/pytensor.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 2.38.1
+Version: 2.38.2
 Summary: Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs.
 Author-email: pymc-devs <pymc.devs@gmail.com>
 License-Expression: BSD-3-Clause