pytensor 3.3.1__tar.gz → 3.3.3__tar.gz

{pytensor-3.3.1/pytensor.egg-info → pytensor-3.3.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 3.3.1
+Version: 3.3.3
 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-3.3.1 → pytensor-3.3.3}/doc/gallery/introduction/pytensor_intro.ipynb

@@ -877,7 +877,7 @@
      "text": [
       "Exp [id A]\n",
       " └─ normal_rv{\"(),()->()\"}.1 [id B]\n",
-      "    ├─ RNG(<Generator(PCG64) at 0x12C30BAE0>) [id C]\n",
+      "    ├─ RNG(Generator(PCG64)) [id C]\n",
       "    ├─ [10] [id D]\n",
       "    ├─ ExpandDims{axis=0} [id E]\n",
       "    │  └─ 0 [id F]\n",
@@ -1495,7 +1495,7 @@
      "output_type": "stream",
      "text": [
       "normal_rv{\"(),()->()\"}.1 [id A] 'z'\n",
-      " ├─ RNG(<Generator(PCG64) at 0x136F5DC40>) [id B]\n",
+      " ├─ RNG(Generator(PCG64)) [id B]\n",
       " ├─ NoneConst{None} [id C]\n",
       " ├─ [0 0] [id D]\n",
       " └─ [1 2] [id E]\n"

{pytensor-3.3.1 → pytensor-3.3.3}/doc/tutorial/prng.rst

@@ -568,7 +568,7 @@ In general, update rules are still respected, but they won't update/rely on the
 >>> jax_fn = pytensor.function([], [x], updates={rng: next_rng}, mode="JAX")
 >>> _ = pytensor.dprint(jax_fn, print_type=True) # doctest: +ELLIPSIS
 uniform_rv{"(),()->()"}.1 [id A] <Scalar(float64, shape=())> 0
- ├─ RNG(<Generator(PCG64) at 0x...>) [id B] <RandomGeneratorType>
+ ├─ RNG(Generator(PCG64)) [id B] <RandomGeneratorType>
  ├─ NoneConst{None} [id C] <NoneTypeT>
  ├─ 0.0 [id D] <Scalar(float32, shape=())>
  └─ 1.0 [id E] <Scalar(float32, shape=())>

{pytensor-3.3.1 → pytensor-3.3.3}/pyproject.toml

@@ -151,6 +151,7 @@ unfixable = [
 
 [tool.ruff.lint.isort]
 lines-after-imports = 2
+known-first-party = ["pytensor", "tests"]
 
 [tool.ruff.lint.per-file-ignores]
 # TODO: Get rid of these:

{pytensor-3.3.1 → pytensor-3.3.3}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2026-09-07T16:04:36+0200",
+ "date": "2026-10-01T22:39:01+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "bb5967ab72cfc29ee9f47c4a2fd0aaeeb1695268",
- "version": "3.3.1"
+ "full-revisionid": "259987437e5efea3678a7d632af90d366c8fd9a2",
+ "version": "3.3.3"
 }
 '''  # END VERSION_JSON
 

{pytensor-3.3.1 → pytensor-3.3.3}/pytensor/d3viz/d3viz.py

@@ -18,7 +18,7 @@ def replace_patterns(x, replace):
 
     Parameters
     ----------
-    s : str
+    x : str
         String on which function is applied
     replace : dict
         `key`, `value` pairs where key is a regular expression and `value` a

{pytensor-3.3.1 → pytensor-3.3.3}/pytensor/graph/fg.py

@@ -321,7 +321,7 @@ class FunctionGraph(AbstractFunctionGraph):
 
         Parameters
         ----------
-        variable : pytensor.graph.basic.Variable
+        var : pytensor.graph.basic.Variable
             The variable to be imported.
         reason : str
             The name of the optimization or operation in progress.

{pytensor-3.3.1 → pytensor-3.3.3}/pytensor/graph/traversal.py

@@ -666,9 +666,9 @@ def toposort_with_orderings(
     Parameters
     ----------
     graphs : list or tuple of Variable instances
-        Graph inputs.
-    outputs : list or tuple of Apply instances
         Graph outputs.
+    blockers : list or tuple of Variable instances
+        Graph inputs.
     orderings : dict
         Keys are `Apply` or `Variable` instances, values are lists of `Apply` or `Variable` instances.
 

{pytensor-3.3.1 → pytensor-3.3.3}/pytensor/link/jax/dispatch/tensor_basic.py

@@ -14,6 +14,7 @@ from pytensor.tensor.basic import (
     Eye,
     Join,
     MakeVector,
+    Nonzero,
     ScalarFromTensor,
     Split,
     TensorFromScalar,
@@ -32,6 +33,13 @@ An example of a graph that can be compiled to JAX:
 """
 
 
+@jax_funcify.register(Nonzero)
+def jax_funcify_Nonzero(op, **kwargs):
+    raise NotImplementedError(
+        "JAX does not support nonzero, its output shape is data-dependent"
+    )
+
+
 @jax_funcify.register(AllocEmpty)
 def jax_funcify_AllocEmpty(op, **kwargs):
     def allocempty(*shape):

{pytensor-3.3.1 → pytensor-3.3.3}/pytensor/link/jax/linker.py

@@ -21,8 +21,8 @@ class JAXLinker(JITLinker):
         # JAX does it his own inplace optimization
         "inplace",
         # There are specific variants for the LU decompositions supported by JAX
-        "reuse_lu_decomposition_multiple_solves",
-        "scan_split_non_sequence_lu_decomposition_solve",
+        "reuse_decomposition_multiple_solves",
+        "scan_split_non_sequence_decomposition_and_solve",
     )
 
     scalar_shape_inputs: tuple[int, ...]

{pytensor-3.3.1 → pytensor-3.3.3}/pytensor/link/jax/ops.py

@@ -10,7 +10,7 @@ from pytensor.compile.mode import Mode
 from pytensor.gradient import DisconnectedType
 from pytensor.graph import Apply, Op, Variable
 from pytensor.tensor.basic import as_tensor, infer_static_shape
-from pytensor.tensor.type import TensorType
+from pytensor.tensor.type import TensorType, discrete_dtypes
 
 
 class JAXOp(Op):
@@ -135,6 +135,13 @@ class JAXOp(Op):
             return outputs[0]
         return outputs
 
+    def connection_pattern(self, node):
+        """Mark discrete inputs as disconnected from every output."""
+        return [
+            [input_type.dtype not in discrete_dtypes] * len(self.output_types)
+            for input_type in self.input_types
+        ]
+
     def pullback(self, inputs, outputs, output_gradients):
         """Compute gradients using JAX's vector-Jacobian product (VJP)."""
         import jax
@@ -146,6 +153,15 @@ class JAXOp(Op):
             if not isinstance(output_grad.type, DisconnectedType)
         ]
 
+        # Integer and boolean inputs are not differentiable. JAX gives them
+        # float0 cotangents, which have no PyTensor equivalent, so they are
+        # held constant in the VJP and reported as disconnected.
+        differentiable_input_indices = [
+            i
+            for i, input_type in enumerate(self.input_types)
+            if input_type.dtype not in discrete_dtypes
+        ]
+
         num_inputs = len(inputs)
 
         def vjp_operation(*args):
@@ -154,15 +170,23 @@ class JAXOp(Op):
             cotangent_vectors = args[num_inputs:]
             assert len(cotangent_vectors) == len(connected_output_indices)
 
-            def restricted_function(*input_values):
-                """Restricted function that only returns connected outputs."""
-                outputs = self.jax_func(*input_values)
+            def restricted_function(*differentiable_values):
+                """Restricted function of the differentiable inputs, returning connected outputs."""
+                all_input_values = list(input_values)
+                for i, value in zip(
+                    differentiable_input_indices, differentiable_values, strict=True
+                ):
+                    all_input_values[i] = value
+                outputs = self.jax_func(*all_input_values)
                 return [
                     outputs[i].astype(self.output_types[i].dtype)
                     for i in connected_output_indices
                 ]
 
-            _primals, vjp_function = jax.vjp(restricted_function, *input_values)
+            _primals, vjp_function = jax.vjp(
+                restricted_function,
+                *(input_values[i] for i in differentiable_input_indices),
+            )
             output_dtypes = [
                 self.output_types[i].dtype for i in connected_output_indices
             ]
@@ -184,15 +208,21 @@ class JAXOp(Op):
         vjp_op = JAXOp(
             self.input_types
             + tuple(self.output_types[i] for i in connected_output_indices),
-            [self.input_types[i] for i in range(num_inputs)],
+            [self.input_types[i] for i in differentiable_input_indices],
             vjp_operation,
             name=name,
         )
 
-        return vjp_op(
+        differentiable_input_gradients = vjp_op(
             *[*inputs, *[output_gradients[i] for i in connected_output_indices]],
             return_list=True,
         )
+        input_gradients = [DisconnectedType()() for _ in inputs]
+        for i, input_gradient in zip(
+            differentiable_input_indices, differentiable_input_gradients, strict=True
+        ):
+            input_gradients[i] = input_gradient
+        return input_gradients
 
 
 def wrap_jax(jax_function=None, *, allow_eval=True):

{pytensor-3.3.1 → pytensor-3.3.3}/pytensor/link/mlx/dispatch/__init__.py

@@ -17,4 +17,5 @@ import pytensor.link.mlx.dispatch.extra_ops
 import pytensor.link.mlx.dispatch.pad
 import pytensor.link.mlx.dispatch.sort
 import pytensor.link.mlx.dispatch.linalg
+import pytensor.link.mlx.dispatch.random
 # isort: on

{pytensor-3.3.1 → pytensor-3.3.3}/pytensor/link/mlx/dispatch/basic.py

@@ -132,7 +132,9 @@ def mlx_typify_tensor(data, dtype=None, **kwargs):
     # and it does so on the CPU too, where float64 is perfectly usable
     if dtype is None and data.dtype == np.float64 and float64_supported():
         dtype = mx.float64
-    return _nan_safe_constant(data, dtype=dtype)
+    # MLX's elementwise kernels misread a non-contiguous buffer, and rewriting
+    # produces such arrays -- `triu` of a transpose, for one
+    return _nan_safe_constant(np.asarray(data, order="C"), dtype=dtype)
 
 
 @mlx_typify.register(slice)

{pytensor-3.3.1 → pytensor-3.3.3}/pytensor/link/mlx/dispatch/blockwise.py

@@ -1,4 +1,5 @@
 import mlx.core as mx
+import numpy as np
 
 from pytensor.link.mlx.dispatch import mlx_funcify
 from pytensor.tensor.blockwise import Blockwise, _check_runtime_broadcast_core
@@ -19,6 +20,29 @@ def funcify_Blockwise(op: Blockwise, node, **kwargs):
     # Hoisted out of the per-call path, unlike Blockwise._check_runtime_broadcast.
     batch_bcast = [inp.type.broadcastable[:batch_ndim] for inp in node.inputs]
 
+    # A core function that sets `natively_batched` runs directly on inputs
+    # broadcast to the common batch shape. MLX's linalg functions take leading
+    # batch dims themselves, and mx.vmap has no rule for LUF or QRF and drops the
+    # triangular flags of solve_triangular.
+    if getattr(core_f, "natively_batched", False):
+
+        def blockwise_native(*args):
+            _check_runtime_broadcast_core(args, batch_bcast, batch_ndim)
+
+            batch_shapes = [
+                arg.shape[: arg.ndim - n_core] for arg, n_core in zip(args, core_ndims)
+            ]
+            batch_shape = np.broadcast_shapes(*batch_shapes)
+            args = [
+                mx.broadcast_to(arg, (*batch_shape, *arg.shape[arg.ndim - n_core :]))
+                for arg, n_core in zip(args, core_ndims)
+            ]
+
+            out = core_f(*args)
+            return tuple(out) if multi_output else out
+
+        return blockwise_native
+
     # Decide batching purely from static shapes so a graph batches identically
     # here and in every other backend: a batch axis broadcasts (is never mapped)
     # only when its static size is exactly 1, or the input lacks it entirely.

{pytensor-3.3.1 → pytensor-3.3.3}/pytensor/link/mlx/dispatch/linalg/decomposition.py

@@ -29,10 +29,9 @@ def mlx_funcify_SVD(op, node, **kwargs):
         )
         return outputs
 
-    if compute_uv:
-        return svd_full
-    else:
-        return svd_S_only
+    svd = svd_full if compute_uv else svd_S_only
+    svd.natively_batched = True
+    return svd
 
 
 @mlx_funcify.register(Cholesky)
@@ -45,6 +44,7 @@ def mlx_funcify_Cholesky(op, node, **kwargs):
             a.astype(dtype=a_dtype, stream=mx.cpu), upper=not lower, stream=mx.cpu
         )
 
+    cholesky.natively_batched = True
     return cholesky
 
 
@@ -70,6 +70,7 @@ def mlx_funcify_LU(op, node, **kwargs):
             U,
         )
 
+    lu.natively_batched = True
     return lu
 
 
@@ -80,6 +81,7 @@ def mlx_funcify_Eig(op, node, **kwargs):
     def eig(x):
         return mx.linalg.eig(x.astype(dtype=X_dtype, stream=mx.cpu), stream=mx.cpu)
 
+    eig.natively_batched = True
     return eig
 
 
@@ -99,6 +101,7 @@ def mlx_funcify_Eigh(op, node, **kwargs):
             a.astype(dtype=X_dtype, stream=mx.cpu), UPLO=UPLO, stream=mx.cpu
         )
 
+    eigh.natively_batched = True
     return eigh
 
 
@@ -118,6 +121,7 @@ def mlx_funcify_Eigvalsh(op, node, **kwargs):
             a.astype(dtype=X_dtype, stream=mx.cpu), UPLO=UPLO, stream=mx.cpu
         )
 
+    eigvalsh.natively_batched = True
     return eigvalsh
 
 
@@ -131,6 +135,7 @@ def mlx_funcify_LUFactor(op, node, **kwargs):
         )
         return lu, pivots.astype(mx.int32, stream=mx.cpu)
 
+    lu_factor.natively_batched = True
     return lu_factor
 
 
@@ -174,4 +179,5 @@ def mlx_funcify_QR(op, node, **kwargs):
             return R
         return Q, R
 
+    qr.natively_batched = True
     return qr

{pytensor-3.3.1 → pytensor-3.3.3}/pytensor/link/mlx/dispatch/linalg/inverse.py

@@ -11,6 +11,7 @@ def mlx_funcify_MatrixInverse(op, node, **kwargs):
     def inv(x):
         return mx.linalg.inv(x.astype(dtype=X_dtype, stream=mx.cpu), stream=mx.cpu)
 
+    inv.natively_batched = True
     return inv
 
 
@@ -21,4 +22,5 @@ def mlx_funcify_MatrixPinv(op, node, **kwargs):
     def pinv(x):
         return mx.linalg.pinv(x.astype(dtype=x_dtype, stream=mx.cpu), stream=mx.cpu)
 
+    pinv.natively_batched = True
     return pinv

{pytensor-3.3.1 → pytensor-3.3.3}/pytensor/link/mlx/dispatch/linalg/solvers.py

@@ -8,6 +8,15 @@ from pytensor.tensor.linalg.solvers.psd import CholeskySolve
 from pytensor.tensor.linalg.solvers.triangular import SolveTriangular
 
 
+def _as_column(b, b_ndim):
+    # MLX treats a 2-d rhs as a matrix, so a batched vector rhs needs an explicit column.
+    return mx.expand_dims(b, -1, stream=mx.cpu) if b_ndim == 1 else b
+
+
+def _from_column(out, b_ndim):
+    return mx.squeeze(out, -1, stream=mx.cpu) if b_ndim == 1 else out
+
+
 @mlx_funcify.register(Solve)
 def mlx_funcify_Solve(op, node, **kwargs):
     assume_a = op.assume_a
@@ -20,30 +29,47 @@ def mlx_funcify_Solve(op, node, **kwargs):
             UserWarning,
         )
 
+    b_ndim = op.b_ndim
+
     def solve(a, b):
-        return mx.linalg.solve(
+        out = mx.linalg.solve(
             a.astype(stream=mx.cpu, dtype=a_dtype),
-            b.astype(stream=mx.cpu, dtype=b_dtype),
+            _as_column(b.astype(stream=mx.cpu, dtype=b_dtype), b_ndim),
             stream=mx.cpu,
         )
+        return _from_column(out, b_ndim)
 
+    solve.natively_batched = True
     return solve
 
 
 @mlx_funcify.register(SolveTriangular)
 def mlx_funcify_SolveTriangular(op, node, **kwargs):
     lower = op.lower
+    unit_diagonal = op.unit_diagonal
     A_dtype = getattr(mx, node.inputs[0].dtype)
     b_dtype = getattr(mx, node.inputs[1].dtype)
+    b_ndim = op.b_ndim
 
     def solve_triangular(A, b):
-        return mx.linalg.solve_triangular(
-            A.astype(stream=mx.cpu, dtype=A_dtype),
-            b.astype(stream=mx.cpu, dtype=b_dtype),
+        A = A.astype(stream=mx.cpu, dtype=A_dtype)
+
+        if unit_diagonal:
+            # MLX's `solve_triangular` has no `unit_diagonal`. LAPACK's `trtrs`
+            # never reads the diagonal in that mode, so overwriting it with ones
+            # gives the same answer.
+            diagonal_mask = mx.eye(A.shape[-1], dtype=mx.bool_, stream=mx.cpu)
+            A = mx.where(diagonal_mask, mx.array(1, dtype=A_dtype), A, stream=mx.cpu)
+
+        out = mx.linalg.solve_triangular(
+            A,
+            _as_column(b.astype(stream=mx.cpu, dtype=b_dtype), b_ndim),
             upper=not lower,
             stream=mx.cpu,
         )
+        return _from_column(out, b_ndim)
 
+    solve_triangular.natively_batched = True
     return solve_triangular
 
 
@@ -52,15 +78,18 @@ def mlx_funcify_CholeskySolve(op, node, **kwargs):
     lower = op.lower
     c_dtype = getattr(mx, node.inputs[0].dtype)
     b_dtype = getattr(mx, node.inputs[1].dtype)
+    b_ndim = op.b_ndim
 
     # MLX has no cho_solve, so with A = L L.T we solve L y = b then L.T x = y.
     def cho_solve(c, b):
         c = c.astype(stream=mx.cpu, dtype=c_dtype)
-        b = b.astype(stream=mx.cpu, dtype=b_dtype)
+        b = _as_column(b.astype(stream=mx.cpu, dtype=b_dtype), b_ndim)
         c_T = mx.swapaxes(c, -1, -2, stream=mx.cpu)
         L, L_T = (c, c_T) if lower else (c_T, c)
 
         y = mx.linalg.solve_triangular(L, b, upper=False, stream=mx.cpu)
-        return mx.linalg.solve_triangular(L_T, y, upper=True, stream=mx.cpu)
+        out = mx.linalg.solve_triangular(L_T, y, upper=True, stream=mx.cpu)
+        return _from_column(out, b_ndim)
 
+    cho_solve.natively_batched = True
     return cho_solve

{pytensor-3.3.1 → pytensor-3.3.3}/pytensor/link/mlx/dispatch/linalg/summary.py

@@ -7,11 +7,11 @@ from pytensor.tensor.linalg.summary import Det, SLogDet
 def _lu_det_parts(x):
     """Compute sign and logdet via LU factorization. Call within a CPU stream context."""
     lu, pivots = mx.linalg.lu_factor(x)
-    diag_u = mx.diagonal(lu)
-    n_swaps = mx.sum(pivots != mx.arange(pivots.shape[0], dtype=pivots.dtype))
+    diag_u = mx.diagonal(lu, axis1=-2, axis2=-1)
+    n_swaps = mx.sum(pivots != mx.arange(pivots.shape[-1], dtype=pivots.dtype), axis=-1)
     pivot_sign = 1 - 2 * (n_swaps % 2)
-    sign = pivot_sign * mx.prod(mx.sign(diag_u))
-    logabsdet = mx.sum(mx.log(mx.abs(diag_u)))
+    sign = pivot_sign * mx.prod(mx.sign(diag_u), axis=-1)
+    logabsdet = mx.sum(mx.log(mx.abs(diag_u)), axis=-1)
     return sign, logabsdet
 
 
@@ -24,6 +24,7 @@ def mlx_funcify_Det(op, node, **kwargs):
             sign, logabsdet = _lu_det_parts(x.astype(dtype=X_dtype))
             return sign * mx.exp(logabsdet)
 
+    det.natively_batched = True
     return det
 
 
@@ -35,4 +36,5 @@ def mlx_funcify_SLogDet(op, node, **kwargs):
         with mx.stream(mx.cpu):
             return _lu_det_parts(x.astype(dtype=X_dtype))
 
+    slogdet.natively_batched = True
     return slogdet

{pytensor-3.3.1 → pytensor-3.3.3}/pytensor/link/mlx/dispatch/pad.py

@@ -1,12 +1,27 @@
 import mlx.core as mx
+import numpy as np
 
 from pytensor.link.mlx.dispatch.basic import mlx_funcify
+from pytensor.link.mlx.dispatch.tensor_basic import mlx_to_list_shape
 from pytensor.tensor.pad import Pad
 
 
+def _runtime_pad_width_pairs(pad_width, ndim):
+    widths = np.array(mlx_to_list_shape(pad_width.flatten())).reshape(pad_width.shape)
+    widths = np.broadcast_to(widths, (ndim, 2))
+    return [tuple(pair) for pair in widths.tolist()]
+
+
 @mlx_funcify.register(Pad)
 def mlx_funcify_pad(op, node, **kwargs):
     pad_mode = op.pad_mode
+    ndim = node.inputs[0].type.ndim
+    static_pairs = None if op.static_pad_width is None else list(op.static_pad_width)
+
+    def pad_width_pairs(pad_width):
+        if static_pairs is not None:
+            return static_pairs
+        return _runtime_pad_width_pairs(pad_width, ndim)
 
     if pad_mode == "constant":
         _, _, constant_values = node.inputs
@@ -18,7 +33,10 @@ def mlx_funcify_pad(op, node, **kwargs):
 
         def constant_pad_fn(x, pad_width, constant_values):
             return mx.pad(
-                x, pad_width, mode="constant", constant_values=constant_values
+                x,
+                pad_width_pairs(pad_width),
+                mode="constant",
+                constant_values=constant_values,
             )
 
         return constant_pad_fn
@@ -26,7 +44,7 @@ def mlx_funcify_pad(op, node, **kwargs):
     elif pad_mode == "edge":
 
         def edge_pad_fn(x, pad_width):
-            return mx.pad(x, pad_width, mode="edge")
+            return mx.pad(x, pad_width_pairs(pad_width), mode="edge")
 
         return edge_pad_fn