pytensor 2.31.4__tar.gz → 2.31.5__tar.gz

{pytensor-2.31.4/pytensor.egg-info → pytensor-2.31.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 2.31.4
+Version: 2.31.5
 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.31.4 → pytensor-2.31.5}/doc/tutorial/gradients.rst

@@ -101,9 +101,12 @@ PyTensor implements the :func:`pytensor.gradient.jacobian` macro that does all
 that is needed to compute the Jacobian. The following text explains how
 to do it manually.
 
+Using Scan
+----------
+
 In order to manually compute the Jacobian of some function ``y`` with
-respect to some parameter ``x`` we need to use `scan`. What we
-do is to loop over the entries in ``y`` and compute the gradient of
+respect to some parameter ``x`` we can use `scan`.
+In this case, we loop over the entries in ``y`` and compute the gradient of
 ``y[i]`` with respect to ``x``.
 
 .. note::
@@ -111,8 +114,7 @@ do is to loop over the entries in ``y`` and compute the gradient of
     `scan` is a generic op in PyTensor that allows writing in a symbolic
     manner all kinds of recurrent equations. While creating
     symbolic loops (and optimizing them for performance) is a hard task,
-    effort is being done for improving the performance of `scan`. We
-    shall return to :ref:`scan<tutloop>` later in this tutorial.
+    efforts are being made to improving the performance of `scan`.
 
 >>> import pytensor
 >>> import pytensor.tensor as pt
@@ -124,9 +126,9 @@ do is to loop over the entries in ``y`` and compute the gradient of
 array([[ 8.,  0.],
        [ 0.,  8.]])
 
-What we do in this code is to generate a sequence of integers from ``0`` to
-``y.shape[0]`` using `pt.arange`. Then we loop through this sequence, and
-at each step, we compute the gradient of element ``y[i]`` with respect to
+This code generates a sequence of integers from ``0`` to
+``y.shape[0]`` using `pt.arange`. Then it loops through this sequence, and
+at each step, computes the gradient of element ``y[i]`` with respect to
 ``x``. `scan` automatically concatenates all these rows, generating a
 matrix which corresponds to the Jacobian.
 
@@ -139,6 +141,31 @@ matrix which corresponds to the Jacobian.
     ``x`` anymore, while ``y[i]`` still is.
 
 
+Using automatic vectorization
+-----------------------------
+An alternative way to build the Jacobian is to vectorize the graph that computes a single row or colum of the jacobian
+We can use `Lop` or `Rop` (more about it below) to obtain the row or column of the jacobian and `vectorize_graph`
+to vectorize it to the full jacobian matrix.
+
+>>> import pytensor
+>>> import pytensor.tensor as pt
+>>> from pytensor.gradient import Lop
+>>> from pytensor.graph import vectorize_graph
+>>> x = pt.dvector('x')
+>>> y = x ** 2
+>>> row_cotangent = pt.dvector("row_cotangent")  # Helper variable, it will be replaced during vectorization
+>>> J_row = Lop(y, x, row_cotangent)
+>>> J = vectorize_graph(J_row, replace={row_cotangent: pt.eye(x.size)})
+>>> f = pytensor.function([x], J)
+>>> f([4, 4])
+array([[ 8.,  0.],
+       [ 0.,  8.]])
+
+This avoids the overhead of scan, at the cost of higher memory usage if the jacobian expression has large intermediate operations.
+Also, not all graphs are safely vectorizable (e.g., if different rows require intermediate operations of different sizes).
+For these reasons `jacobian` uses scan by default. The behavior can be changed by setting `vectorize=True`.
+
+
 Computing the Hessian
 =====================
 

{pytensor-2.31.4 → pytensor-2.31.5}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2025-06-24T11:04:32+0200",
+ "date": "2025-06-30T17:01:33+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "46c2341072820cc1cebe6f77fe754ed63a703988",
- "version": "2.31.4"
+ "full-revisionid": "2ab60b39e4622b1d9d677ab3e598d752bd25b17e",
+ "version": "2.31.5"
 }
 '''  # END VERSION_JSON
 

{pytensor-2.31.4 → pytensor-2.31.5}/pytensor/gradient.py

@@ -11,7 +11,7 @@ import numpy as np
 import pytensor
 from pytensor.compile.ops import ViewOp
 from pytensor.configdefaults import config
-from pytensor.graph import utils
+from pytensor.graph import utils, vectorize_graph
 from pytensor.graph.basic import Apply, NominalVariable, Variable
 from pytensor.graph.null_type import NullType, null_type
 from pytensor.graph.op import get_test_values
@@ -703,15 +703,15 @@ def grad(
         grad_dict[var] = g_var
 
     def handle_disconnected(var):
-        message = (
-            "grad method was asked to compute the gradient "
-            "with respect to a variable that is not part of "
-            "the computational graph of the cost, or is used "
-            f"only by a non-differentiable operator: {var}"
-        )
         if disconnected_inputs == "ignore":
-            pass
+            return
         elif disconnected_inputs == "warn":
+            message = (
+                "grad method was asked to compute the gradient "
+                "with respect to a variable that is not part of "
+                "the computational graph of the cost, or is used "
+                f"only by a non-differentiable operator: {var}"
+            )
             warnings.warn(message, stacklevel=2)
         elif disconnected_inputs == "raise":
             message = utils.get_variable_trace_string(var)
@@ -2021,13 +2021,19 @@ GradientError: numeric gradient and analytic gradient exceed tolerance:
 Exception args: {args_msg}"""
 
 
-def jacobian(expression, wrt, consider_constant=None, disconnected_inputs="raise"):
+def jacobian(
+    expression,
+    wrt,
+    consider_constant=None,
+    disconnected_inputs="raise",
+    vectorize=False,
+):
     """
     Compute the full Jacobian, row by row.
 
     Parameters
     ----------
-    expression : Vector (1-dimensional) :class:`~pytensor.graph.basic.Variable`
+    expression :class:`~pytensor.graph.basic.Variable`
         Values that we are differentiating (that we want the Jacobian of)
     wrt : :class:`~pytensor.graph.basic.Variable` or list of Variables
         Term[s] with respect to which we compute the Jacobian
@@ -2051,18 +2057,18 @@ def jacobian(expression, wrt, consider_constant=None, disconnected_inputs="raise
         output, then a zero variable is returned. The return value is
         of same type as `wrt`: a list/tuple or TensorVariable in all cases.
     """
+    from pytensor.tensor.basic import eye
+    from pytensor.tensor.extra_ops import broadcast_to
 
     if not isinstance(expression, Variable):
         raise TypeError("jacobian expects a Variable as `expression`")
 
-    if expression.ndim > 1:
-        raise ValueError(
-            "jacobian expects a 1 dimensional variable as `expression`."
-            " If not use flatten to make it a vector"
-        )
-
     using_list = isinstance(wrt, list)
     using_tuple = isinstance(wrt, tuple)
+    grad_kwargs = {
+        "consider_constant": consider_constant,
+        "disconnected_inputs": disconnected_inputs,
+    }
 
     if isinstance(wrt, list | tuple):
         wrt = list(wrt)
@@ -2070,43 +2076,55 @@ def jacobian(expression, wrt, consider_constant=None, disconnected_inputs="raise
         wrt = [wrt]
 
     if all(expression.type.broadcastable):
-        # expression is just a scalar, use grad
-        return as_list_or_tuple(
-            using_list,
-            using_tuple,
-            grad(
-                expression.squeeze(),
-                wrt,
-                consider_constant=consider_constant,
-                disconnected_inputs=disconnected_inputs,
-            ),
+        jacobian_matrices = grad(expression.squeeze(), wrt, **grad_kwargs)
+
+    elif vectorize:
+        expression_flat = expression.ravel()
+        row_tangent = _float_ones_like(expression_flat).type("row_tangent")
+        jacobian_single_rows = Lop(expression.ravel(), wrt, row_tangent, **grad_kwargs)
+
+        n_rows = expression_flat.size
+        jacobian_matrices = vectorize_graph(
+            jacobian_single_rows,
+            replace={row_tangent: eye(n_rows, dtype=row_tangent.dtype)},
         )
+        if disconnected_inputs != "raise":
+            # If the input is disconnected from the cost, `vectorize_graph` has no effect on the respective jacobian
+            # We have to broadcast the zeros explicitly here
+            for i, (jacobian_single_row, jacobian_matrix) in enumerate(
+                zip(jacobian_single_rows, jacobian_matrices, strict=True)
+            ):
+                if jacobian_single_row.ndim == jacobian_matrix.ndim:
+                    jacobian_matrices[i] = broadcast_to(
+                        jacobian_matrix, shape=(n_rows, *jacobian_matrix.shape)
+                    )
 
-    def inner_function(*args):
-        idx = args[0]
-        expr = args[1]
-        rvals = []
-        for inp in args[2:]:
-            rval = grad(
-                expr[idx],
-                inp,
-                consider_constant=consider_constant,
-                disconnected_inputs=disconnected_inputs,
+    else:
+
+        def inner_function(*args):
+            idx, expr, *wrt = args
+            return grad(expr[idx], wrt, **grad_kwargs)
+
+        jacobian_matrices, updates = pytensor.scan(
+            inner_function,
+            sequences=pytensor.tensor.arange(expression.size),
+            non_sequences=[expression.ravel(), *wrt],
+            return_list=True,
+        )
+        if updates:
+            raise ValueError(
+                "The scan used to build the jacobian matrices returned a list of updates"
             )
-            rvals.append(rval)
-        return rvals
-
-    # Computing the gradients does not affect the random seeds on any random
-    # generator used n expression (because during computing gradients we are
-    # just backtracking over old values. (rp Jan 2012 - if anyone has a
-    # counter example please show me)
-    jacobs, updates = pytensor.scan(
-        inner_function,
-        sequences=pytensor.tensor.arange(expression.shape[0]),
-        non_sequences=[expression, *wrt],
-    )
-    assert not updates, "Scan has returned a list of updates; this should not happen."
-    return as_list_or_tuple(using_list, using_tuple, jacobs)
+
+    if jacobian_matrices[0].ndim < (expression.ndim + wrt[0].ndim):
+        # There was some raveling or squeezing done prior to getting the jacobians
+        # Reshape into original shapes
+        jacobian_matrices = [
+            jac_matrix.reshape((*expression.shape, *w.shape))
+            for jac_matrix, w in zip(jacobian_matrices, wrt, strict=True)
+        ]
+
+    return as_list_or_tuple(using_list, using_tuple, jacobian_matrices)
 
 
 def hessian(cost, wrt, consider_constant=None, disconnected_inputs="raise"):

{pytensor-2.31.4 → pytensor-2.31.5}/pytensor/graph/replace.py

@@ -232,13 +232,13 @@ def vectorize_graph(
 def vectorize_graph(
     outputs: Sequence[Variable],
     replace: Mapping[Variable, Variable],
-) -> Sequence[Variable]: ...
+) -> list[Variable]: ...
 
 
 def vectorize_graph(
     outputs: Variable | Sequence[Variable],
     replace: Mapping[Variable, Variable],
-) -> Variable | Sequence[Variable]:
+) -> Variable | list[Variable]:
     """Vectorize outputs graph given mapping from old variables to expanded counterparts version.
 
     Expanded dimensions must be on the left. Behavior is similar to the functional `numpy.vectorize`.

{pytensor-2.31.4 → pytensor-2.31.5}/pytensor/tensor/basic.py

@@ -3081,6 +3081,10 @@ def flatten(x, ndim=1):
     else:
         dims = (-1,)
 
+    if len(dims) == _x.ndim:
+        # Nothing to ravel
+        return _x
+
     x_reshaped = _x.reshape(dims)
     shape_kept_dims = _x.type.shape[: ndim - 1]
     bcast_new_dim = builtins.all(s == 1 for s in _x.type.shape[ndim - 1 :])

{pytensor-2.31.4 → pytensor-2.31.5}/pytensor/tensor/random/op.py

@@ -392,6 +392,13 @@ class RandomVariable(RNGConsumerOp):
         out_type = TensorType(dtype=self.dtype, shape=static_shape)
         outputs = (rng.type(), out_type())
 
+        if self.dtype == "floatX":
+            # Commit to a specific float type if the Op is still using "floatX"
+            dtype = config.floatX
+            props = self._props_dict()
+            props["dtype"] = dtype
+            self = type(self)(**props)
+
         return Apply(self, inputs, outputs)
 
     def batch_ndim(self, node: Apply) -> int:

{pytensor-2.31.4 → pytensor-2.31.5}/pytensor/tensor/rewriting/linalg.py

@@ -47,8 +47,10 @@ from pytensor.tensor.rewriting.basic import (
 from pytensor.tensor.slinalg import (
     BlockDiagonal,
     Cholesky,
+    CholeskySolve,
     Solve,
     SolveBase,
+    SolveTriangular,
     _bilinear_solve_discrete_lyapunov,
     block_diag,
     cholesky,
@@ -908,6 +910,11 @@ def rewrite_cholesky_diag_to_sqrt_diag(fgraph, node):
         return None
 
     [input] = node.inputs
+
+    # Check if input is a (1, 1) matrix
+    if all(input.type.broadcastable[-2:]):
+        return [pt.sqrt(input)]
+
     # Check for use of pt.diag first
     if (
         input.owner
@@ -1020,3 +1027,42 @@ def slogdet_specialization(fgraph, node):
             k: slogdet_specialization_map[v] for k, v in dummy_replacements.items()
         }
         return replacements
+
+
+@register_stabilize
+@register_canonicalize
+@node_rewriter([Blockwise])
+def scalar_solve_to_division(fgraph, node):
+    """
+    Replace solve(a, b) with b / a if a is a (1, 1) matrix
+    """
+
+    core_op = node.op.core_op
+    if not isinstance(core_op, SolveBase):
+        return None
+
+    a, b = node.inputs
+    old_out = node.outputs[0]
+    if not all(a.broadcastable[-2:]):
+        return None
+
+    # Special handling for different types of solve
+    match core_op:
+        case SolveTriangular():
+            # Corner case: if user asked for a triangular solve with a unit diagonal, a is taken to be 1
+            new_out = b / a if not core_op.unit_diagonal else b
+        case CholeskySolve():
+            new_out = b / a**2
+        case Solve():
+            new_out = b / a
+        case _:
+            raise NotImplementedError(
+                f"Unsupported core_op type: {type(core_op)} in scalar_solve_to_divison"
+            )
+
+    if core_op.b_ndim == 1:
+        new_out = new_out.squeeze(-1)
+
+    copy_stack_trace(old_out, new_out)
+
+    return [new_out]

{pytensor-2.31.4 → pytensor-2.31.5}/pytensor/tensor/rewriting/math.py

@@ -400,6 +400,37 @@ def local_exp_log(fgraph, node):
         return [exp(x)]
 
 
+@register_canonicalize
+@register_specialize
+@node_rewriter([sqrt, sqr])
+def local_sqrt_sqr(fgraph, node):
+    x = node.inputs[0]
+
+    if not (x.owner and isinstance(x.owner.op, Elemwise)):
+        return
+
+    prev_op = x.owner.op.scalar_op
+    node_op = node.op.scalar_op
+
+    # Case for sqrt(sqr(x)) -> |x|
+    if isinstance(prev_op, ps.Sqrt) and isinstance(node_op, ps.Sqr):
+        new_out = pt_abs(x.owner.inputs[0])
+        old_out = node.outputs[0]
+
+        # Handle potential integer to float cast by sqr
+        if new_out.dtype != old_out.dtype:
+            new_out = cast(new_out, old_out.dtype)
+        return [new_out]
+
+    # Case for sqr(sqrt(x)) -> x
+    if isinstance(prev_op, ps.Sqr) and isinstance(node_op, ps.Sqrt):
+        x = x.owner.inputs[0]
+        old_out = node.outputs[0]
+        new_out = switch(ge(x, 0), x, np.asarray(np.nan, old_out.dtype))
+
+        return [new_out]
+
+
 @register_specialize
 @node_rewriter([exp, expm1])
 def local_exp_log_nan_switch(fgraph, node):

{pytensor-2.31.4 → pytensor-2.31.5}/pytensor/xtensor/__init__.py

@@ -1,7 +1,7 @@
 import warnings
 
 import pytensor.xtensor.rewriting
-from pytensor.xtensor import linalg
+from pytensor.xtensor import linalg, random
 from pytensor.xtensor.math import dot
 from pytensor.xtensor.shape import concat
 from pytensor.xtensor.type import (

{pytensor-2.31.4 → pytensor-2.31.5}/pytensor/xtensor/math.py

@@ -101,7 +101,7 @@ round = _as_xelemwise(ps.round_half_to_even)
 maximum = _as_xelemwise(ps.scalar_maximum)
 minimum = _as_xelemwise(ps.scalar_minimum)
 second = _as_xelemwise(ps.second)
-sigmoid = _as_xelemwise(ps.sigmoid)
+sigmoid = expit = _as_xelemwise(ps.sigmoid)
 sign = _as_xelemwise(ps.sign)
 sin = _as_xelemwise(ps.sin)
 sinh = _as_xelemwise(ps.sinh)

{pytensor-2.31.4 → pytensor-2.31.5}/pytensor/xtensor/random.py

@@ -5,8 +5,8 @@ from typing import Literal
 import pytensor.tensor.random.basic as ptr
 from pytensor.graph.basic import Variable
 from pytensor.tensor.random.op import RandomVariable
-from pytensor.xtensor import as_xtensor
 from pytensor.xtensor.math import sqrt
+from pytensor.xtensor.type import as_xtensor
 from pytensor.xtensor.vectorization import XRV
 
 
@@ -14,6 +14,7 @@ def _as_xrv(
     core_op: RandomVariable,
     core_inps_dims_map: Sequence[Sequence[int]] | None = None,
     core_out_dims_map: Sequence[int] | None = None,
+    name: str | None = None,
 ):
     """Helper function to define an XRV constructor.
 
@@ -41,7 +42,14 @@ def _as_xrv(
         core_out_dims_map = tuple(range(core_op.ndim_supp))
 
     core_dims_needed = max(
-        (*(len(i) for i in core_inps_dims_map), len(core_out_dims_map)), default=0
+        max(
+            (
+                max((entry + 1 for entry in dims_map), default=0)
+                for dims_map in core_inps_dims_map
+            ),
+            default=0,
+        ),
+        max((entry + 1 for entry in core_out_dims_map), default=0),
     )
 
     @wraps(core_op)
@@ -76,7 +84,10 @@ def _as_xrv(
             extra_dims = {}
 
         return XRV(
-            core_op, core_dims=full_core_dims, extra_dims=tuple(extra_dims.keys())
+            core_op,
+            core_dims=full_core_dims,
+            extra_dims=tuple(extra_dims.keys()),
+            name=name,
         )(rng, *extra_dims.values(), *params)
 
     return xrv_constructor

{pytensor-2.31.4 → pytensor-2.31.5}/pytensor/xtensor/rewriting/utils.py

@@ -1,6 +1,7 @@
 from pytensor.compile import optdb
-from pytensor.graph.rewriting.basic import NodeRewriter
+from pytensor.graph.rewriting.basic import NodeRewriter, in2out
 from pytensor.graph.rewriting.db import EquilibriumDB, RewriteDatabase
+from pytensor.tensor.rewriting.ofg import inline_ofg_expansion
 
 
 lower_xtensor_db = EquilibriumDB(ignore_newtrees=False)
@@ -14,6 +15,15 @@ optdb.register(
     position=0.1,
 )
 
+# Register OFG inline again after lowering xtensor
+optdb.register(
+    "inline_ofg_expansion_xtensor",
+    in2out(inline_ofg_expansion),
+    "fast_run",
+    "fast_compile",
+    position=0.11,
+)
+
 
 def register_lower_xtensor(
     node_rewriter: RewriteDatabase | NodeRewriter | str, *tags: str, **kwargs

{pytensor-2.31.4 → pytensor-2.31.5}/pytensor/xtensor/rewriting/vectorization.py

@@ -116,7 +116,7 @@ def lower_rv(fgraph, node):
         size = [*extra_dim_lengths, *param_batch_shape]
 
     # RVs are their own core Op
-    new_next_rng, tensor_out = core_op(*tensor_params, rng=rng, size=size).owner.outputs
+    new_next_rng, tensor_out = core_op.make_node(rng, size, *tensor_params).outputs
 
     # Convert output Tensors to XTensors
     new_out = xtensor_from_tensor(tensor_out, dims=old_out.type.dims)

{pytensor-2.31.4 → pytensor-2.31.5}/pytensor/xtensor/shape.py

@@ -189,7 +189,7 @@ class Transpose(XOp):
 
 def transpose(
     x,
-    *dims: str | EllipsisType,
+    *dim: str | EllipsisType,
     missing_dims: Literal["raise", "warn", "ignore"] = "raise",
 ):
     """Transpose dimensions of the tensor.
@@ -198,7 +198,7 @@ def transpose(
     ----------
     x : XTensorVariable
         Input tensor to transpose.
-    *dims : str
+    *dim : str
         Dimensions to transpose to. Can include ellipsis (...) to represent
         remaining dimensions in their original order.
     missing_dims : {"raise", "warn", "ignore"}, optional
@@ -220,7 +220,7 @@ def transpose(
     # Validate dimensions
     x = as_xtensor(x)
     x_dims = x.type.dims
-    invalid_dims = set(dims) - {..., *x_dims}
+    invalid_dims = set(dim) - {..., *x_dims}
     if invalid_dims:
         if missing_dims != "ignore":
             msg = f"Dimensions {invalid_dims} do not exist. Expected one or more of: {x_dims}"
@@ -229,21 +229,27 @@ def transpose(
             else:
                 warnings.warn(msg)
         # Handle missing dimensions if not raising
-        dims = tuple(d for d in dims if d in x_dims or d is ...)
-
-    if dims == () or dims == (...,):
-        dims = tuple(reversed(x_dims))
-    elif ... in dims:
-        if dims.count(...) > 1:
+        dim = tuple(d for d in dim if d in x_dims or d is ...)
+
+    if dim == ():
+        dim = tuple(reversed(x_dims))
+    elif dim == (...,):
+        dim = x_dims
+    elif ... in dim:
+        if dim.count(...) > 1:
             raise ValueError("Ellipsis (...) can only appear once in the dimensions")
         # Handle ellipsis expansion
-        ellipsis_idx = dims.index(...)
-        pre = dims[:ellipsis_idx]
-        post = dims[ellipsis_idx + 1 :]
+        ellipsis_idx = dim.index(...)
+        pre = dim[:ellipsis_idx]
+        post = dim[ellipsis_idx + 1 :]
         middle = [d for d in x_dims if d not in pre + post]
-        dims = (*pre, *middle, *post)
+        dim = (*pre, *middle, *post)
+
+    if dim == x_dims:
+        # No-op transpose
+        return x
 
-    return Transpose(typing.cast(tuple[str], dims))(x)
+    return Transpose(dims=typing.cast(tuple[str], dim))(x)
 
 
 class Concat(XOp):

{pytensor-2.31.4 → pytensor-2.31.5}/pytensor/xtensor/type.py

@@ -71,6 +71,8 @@ class XTensorType(Type, HasDataType, HasShape):
         self.name = name
         self.numpy_dtype = np.dtype(self.dtype)
         self.filter_checks_isfinite = False
+        # broadcastable is here just for code that would work fine with XTensorType but checks for it
+        self.broadcastable = (False,) * self.ndim
 
     def clone(
         self,
@@ -93,6 +95,10 @@ class XTensorType(Type, HasDataType, HasShape):
             self, value, strict=strict, allow_downcast=allow_downcast
         )
 
+    @staticmethod
+    def may_share_memory(a, b):
+        return TensorType.may_share_memory(a, b)
+
     def filter_variable(self, other, allow_convert=True):
         if not isinstance(other, Variable):
             # The value is not a Variable: we cast it into
@@ -160,7 +166,7 @@ class XTensorType(Type, HasDataType, HasShape):
         return None
 
     def __repr__(self):
-        return f"XTensorType({self.dtype}, {self.dims}, {self.shape})"
+        return f"XTensorType({self.dtype}, shape={self.shape}, dims={self.dims})"
 
     def __hash__(self):
         return hash((type(self), self.dtype, self.shape, self.dims))
@@ -691,14 +697,14 @@ class XTensorVariable(Variable[_XTensorTypeType, OptionalApplyType]):
     # https://docs.xarray.dev/en/latest/api.html#id8
     def transpose(
         self,
-        *dims: str | EllipsisType,
+        *dim: str | EllipsisType,
         missing_dims: Literal["raise", "warn", "ignore"] = "raise",
     ):
         """Transpose dimensions of the tensor.
 
         Parameters
         ----------
-        *dims : str | Ellipsis
+        *dim : str | Ellipsis
             Dimensions to transpose. If empty, performs a full transpose.
             Can use ellipsis (...) to represent remaining dimensions.
         missing_dims : {"raise", "warn", "ignore"}, default="raise"
@@ -718,7 +724,7 @@ class XTensorVariable(Variable[_XTensorTypeType, OptionalApplyType]):
             If missing_dims="raise" and any dimensions don't exist.
             If multiple ellipsis are provided.
         """
-        return px.shape.transpose(self, *dims, missing_dims=missing_dims)
+        return px.shape.transpose(self, *dim, missing_dims=missing_dims)
 
     def stack(self, dim, **dims):
         return px.shape.stack(self, dim, **dims)

{pytensor-2.31.4 → pytensor-2.31.5}/pytensor/xtensor/vectorization.py

@@ -142,8 +142,12 @@ class XRV(XOp, RNGConsumerOp):
         core_op,
         core_dims: tuple[tuple[tuple[str, ...], ...], tuple[str, ...]],
         extra_dims: tuple[str, ...],
+        name: str | None = None,
     ):
         super().__init__()
+        if name is None:
+            name = getattr(core_op, "name", None)
+        self.name = name
         self.core_op = core_op
         inps_core_dims, out_core_dims = core_dims
         for operand_dims in (*inps_core_dims, out_core_dims):
@@ -154,6 +158,15 @@ class XRV(XOp, RNGConsumerOp):
             raise ValueError("size_dims must be unique")
         self.extra_dims = tuple(extra_dims)
 
+    def __str__(self):
+        if self.name is not None:
+            name = self.name
+            attrs = f"(core_dims={self.core_dims}, extra_dims={self.extra_dims})"
+        else:
+            name = self.__class__.__name__
+            attrs = f"(core_op={self.core_op}, core_dims={self.core_dims}, extra_dims={self.extra_dims})"
+        return f"{name}({attrs})"
+
     def update(self, node):
         # RNG input and update are the first input and output respectively
         return {node.inputs[0]: node.outputs[0]}

{pytensor-2.31.4 → pytensor-2.31.5/pytensor.egg-info}/PKG-INFO

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