pytensor 3.3.2__tar.gz → 3.3.3__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytensor
-Version: 3.3.2
+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.2 → 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.2 → 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.2 → pytensor-3.3.3}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2026-09-14T16:15:54+0200",
+ "date": "2026-10-01T22:39:01+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "895f7795e2ae8ea16e3a4097ce4b92c61c8faf76",
- "version": "3.3.2"
+ "full-revisionid": "259987437e5efea3678a7d632af90d366c8fd9a2",
+ "version": "3.3.3"
 }
 '''  # END VERSION_JSON
 

{pytensor-3.3.2 → 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.2 → 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.2 → 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.2 → 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.2 → 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.2 → 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.2 → 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,13 +29,17 @@ 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
 
 
@@ -36,6 +49,7 @@ def mlx_funcify_SolveTriangular(op, node, **kwargs):
     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):
         A = A.astype(stream=mx.cpu, dtype=A_dtype)
@@ -47,13 +61,15 @@ def mlx_funcify_SolveTriangular(op, node, **kwargs):
             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)
 
-        return mx.linalg.solve_triangular(
+        out = mx.linalg.solve_triangular(
             A,
-            b.astype(stream=mx.cpu, dtype=b_dtype),
+            _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
 
 
@@ -62,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.2 → 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.2 → 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
 

{pytensor-3.3.2 → pytensor-3.3.3}/pytensor/link/mlx/dispatch/shape.py

@@ -1,6 +1,8 @@
 import mlx.core as mx
 
+from pytensor.graph import Constant
 from pytensor.link.mlx.dispatch.basic import mlx_funcify
+from pytensor.link.mlx.dispatch.tensor_basic import mlx_to_list_shape
 from pytensor.tensor.shape import Reshape, Shape, Shape_i, SpecifyShape
 
 
@@ -37,8 +39,18 @@ def mlx_funcify_Shape_i(op, node, **kwargs):
 
 
 @mlx_funcify.register(Reshape)
-def mlx_funcify_Reshape(op, **kwargs):
-    def reshape(x, shp):
-        return mx.reshape(x, shp)
+def mlx_funcify_Reshape(op, node, **kwargs):
+    shape = node.inputs[1]
+
+    if isinstance(shape, Constant):
+        constant_shape = tuple(int(dim) for dim in shape.data)
+
+        def reshape(x, shp):
+            return mx.reshape(x, constant_shape)
+
+    else:
+
+        def reshape(x, shp):
+            return mx.reshape(x, mlx_to_list_shape(shp))
 
     return reshape

pytensor-3.3.3/pytensor/link/mlx/dispatch/sort.py

@@ -0,0 +1,52 @@
+import warnings
+
+import mlx.core as mx
+
+from pytensor.link.mlx.dispatch.basic import convert_dtype_to_mlx, mlx_funcify
+from pytensor.link.mlx.dispatch.tensor_basic import coerce_to_int
+from pytensor.tensor.basic import get_scalar_constant_value
+from pytensor.tensor.exceptions import NotScalarConstantError
+from pytensor.tensor.sort import ArgSortOp, SortOp
+
+
+def _warn_unsupported_kind(op, name):
+    if op.kind != "quicksort":
+        warnings.warn(
+            message=f"MLX {name} does not support the kind argument (got kind={op.kind}). "
+            "The argument will be ignored.",
+            category=UserWarning,
+        )
+
+
+def _static_axis(node):
+    try:
+        return int(get_scalar_constant_value(node.inputs[1]))
+    except NotScalarConstantError:
+        return None
+
+
+def _resolve_axis(static_axis, axis):
+    return coerce_to_int(axis) if static_axis is None else static_axis
+
+
+@mlx_funcify.register(SortOp)
+def mlx_funcify_Sort(op, node, **kwargs):
+    _warn_unsupported_kind(op, "sort")
+    static_axis = _static_axis(node)
+
+    def sort(x, axis):
+        return mx.sort(x, axis=_resolve_axis(static_axis, axis))
+
+    return sort
+
+
+@mlx_funcify.register(ArgSortOp)
+def mlx_funcify_ArgSort(op, node, **kwargs):
+    _warn_unsupported_kind(op, "argsort")
+    static_axis = _static_axis(node)
+    out_dtype = convert_dtype_to_mlx(node.outputs[0].dtype)
+
+    def argsort(x, axis):
+        return mx.argsort(x, axis=_resolve_axis(static_axis, axis)).astype(out_dtype)
+
+    return argsort

{pytensor-3.3.2 → pytensor-3.3.3}/pytensor/link/mlx/dispatch/subtensor.py

@@ -10,6 +10,21 @@ from pytensor.tensor.subtensor import (
 )
 
 
+def _with_int_slice_bounds(indices):
+    # Slice bounds arrive as 0-d arrays, which MLX slices reject.
+    return tuple(
+        slice(
+            *(
+                None if bound is None else int(bound)
+                for bound in (index.start, index.stop, index.step)
+            )
+        )
+        if isinstance(index, slice)
+        else index
+        for index in indices
+    )
+
+
 @mlx_funcify.register(Subtensor)
 def mlx_funcify_Subtensor(op, node, **kwargs):
     def subtensor(x, *ilists):
@@ -27,7 +42,7 @@ def mlx_funcify_Subtensor(op, node, **kwargs):
 @mlx_funcify.register(AdvancedSubtensor)
 def mlx_funcify_AdvancedSubtensor(op, node, **kwargs):
     def advanced_subtensor(x, *ilists):
-        indices = indices_from_subtensor(ilists, op.idx_list)
+        indices = _with_int_slice_bounds(indices_from_subtensor(ilists, op.idx_list))
         if len(indices) == 1:
             indices = indices[0]
 
@@ -96,9 +111,12 @@ def mlx_funcify_AdvancedIncSubtensor(op, node, **kwargs):
         def mlx_fn(x, indices, y):
             return x.at[indices].add(y)
 
-    def advancedincsubtensor(x, y, *ilist, mlx_fn=mlx_fn):
+    def advancedincsubtensor(x, y, *ilist, mlx_fn=mlx_fn, idx_list=op.idx_list):
         op._check_runtime_broadcast_of_vector_index(node, x, y, ilist[0])
 
-        return mlx_fn(x, ilist, y)
+        # mirrors AdvancedIncSubtensor.perform
+        indices = _with_int_slice_bounds(indices_from_subtensor(ilist, idx_list))
+
+        return mlx_fn(x, indices, y)
 
     return advancedincsubtensor

{pytensor-3.3.2 → pytensor-3.3.3}/pytensor/link/mlx/dispatch/tensor_basic.py

@@ -19,14 +19,14 @@ from pytensor.tensor.basic import (
 from pytensor.tensor.exceptions import NotScalarConstantError
 
 
-MLX_DYNAMIC_SHAPE_ERROR = (
-    "MLX compilation limitation: Alloc operations with dynamic shapes "
-    "cannot be used inside compiled functions. This is because MLX "
-    "compilation forbids evaluating arrays to extract shape values. "
+MLX_DYNAMIC_VALUE_ERROR = (
+    "MLX compilation limitation: integer arguments (shapes, axes, pad widths) "
+    "computed from symbolic inputs cannot be used inside compiled functions. "
+    "This is because MLX compilation forbids evaluating arrays to extract "
+    "their values. "
     "\n\nWorkarounds:"
-    "\n1. Avoid using Alloc with dynamic shapes in compiled contexts"
-    "\n2. Use static shapes when possible"
-    "\n3. Move Alloc operations outside compiled functions"
+    "\n1. Use constant values for these arguments when possible"
+    "\n2. Compile with `MLXLinker(use_compile=False)`"
 )
 
 
@@ -147,7 +147,7 @@ def mlx_funcify_AllocEmpty(op, node, **kwargs):
         resolved_shape = (
             _resolve_shape(static_dims, shape)
             if static_dims is not None
-            else tuple(_coerce_to_int(dim) for dim in shape)
+            else tuple(mlx_to_list_shape(shape))
         )
         return mx.zeros(resolved_shape, dtype=dtype)
 
@@ -167,7 +167,7 @@ def mlx_funcify_Alloc(op, node, **kwargs):
         resolved_shape = (
             _resolve_shape(static_dims, shape)
             if static_dims is not None
-            else tuple(_coerce_to_int(dim) for dim in shape)
+            else tuple(mlx_to_list_shape(shape))
         )
         result = mx.broadcast_to(x, resolved_shape)
         if node_inputs is not None:
@@ -222,12 +222,12 @@ def _resolve_shape(static_dims, runtime_shape):
 
     resolved = []
     for const_dim, dim in zip(static_dims, runtime_shape, strict=True):
-        resolved.append(const_dim if const_dim is not None else _coerce_to_int(dim))
+        resolved.append(const_dim if const_dim is not None else coerce_to_int(dim))
 
     return tuple(resolved)
 
 
-def _coerce_to_int(value):
+def coerce_to_int(value):
     if isinstance(value, np.integer | int):
         return int(value)
     try:
@@ -235,21 +235,22 @@ def _coerce_to_int(value):
             return int(value.item())
         return int(value)
     except (ValueError, TypeError) as exc:
-        _rethrow_dynamic_shape_error(exc)
+        _rethrow_dynamic_value_error(exc)
         raise
 
 
 def mlx_to_list_shape(size) -> list[int]:
-    """Convert a size value (mx.array, np.ndarray, or sequence) to a plain Python list of ints.
-
-    Used by random variable dispatch to normalise the ``size`` argument, which
-    PyTensor may pass as an ``mx.array`` or ``np.ndarray`` rather than a plain
-    Python list.
-    """
-    return [_coerce_to_int(x) for x in size]
+    """Convert a shape given as an ``mx.array``, ``np.ndarray`` or sequence to a list of Python ints."""
+    if isinstance(size, mx.array):
+        try:
+            return [int(dim) for dim in size.tolist()]
+        except (ValueError, TypeError) as exc:
+            _rethrow_dynamic_value_error(exc)
+            raise
+    return [coerce_to_int(x) for x in size]
 
 
-def _rethrow_dynamic_shape_error(exc):
+def _rethrow_dynamic_value_error(exc):
     msg = str(exc)
     if "[eval] Attempting to eval an array during function transformations" in msg:
-        raise ValueError(f"{MLX_DYNAMIC_SHAPE_ERROR}\n\nOriginal error: {msg}") from exc
+        raise ValueError(f"{MLX_DYNAMIC_VALUE_ERROR}\n\nOriginal error: {msg}") from exc