Feedback from Ricardo

Author: cetagostini

pytensor/link/mlx/dispatch/basic.py

@@ -7,6 +7,7 @@
 import numpy as np
 
 from pytensor.compile.ops import DeepCopyOp
+from pytensor.graph import Constant
 from pytensor.graph.fg import FunctionGraph
 from pytensor.link.utils import fgraph_to_python
 from pytensor.raise_op import Assert, CheckAndRaise
@@ -24,7 +25,6 @@ def mlx_typify_tensor(data, dtype=None, **kwargs):
 
 @mlx_typify.register(slice)
 @mlx_typify.register(NoneType)
-@mlx_typify.register(np.number)
 @mlx_typify.register(mx.array)
 def mlx_typify_no_conversion_needed(data, **kwargs):
     return data
@@ -36,6 +36,19 @@ def mlx_typify_python_scalar(data, **kwargs):
     return mx.array(data)
 
 
+@mlx_typify.register(bool)
+@mlx_typify.register(np.bool_)
+def mlx_typify_bool(data, **kwargs):
+    return bool(data)
+
+
+@mlx_typify.register(np.integer)
+@mlx_typify.register(np.floating)
+@mlx_typify.register(np.complexfloating)
+def mlx_typify_numpy_scalar(data, **kwargs):
+    return mx.array(data)
+
+
 @singledispatch
 def mlx_funcify(op, node=None, storage_map=None, **kwargs):
     """Create a MLX compatible function from an PyTensor `Op`."""
@@ -72,9 +85,13 @@ def deepcopyop(x):
 
 @mlx_funcify.register(Assert)
 @mlx_funcify.register(CheckAndRaise)
-def mlx_funcify_CheckAndRaise(op, **kwargs):
+def mlx_funcify_CheckAndRaise(op, node, **kwargs):
+    conds = node.inputs[1:]
+    if any(isinstance(cond, Constant) and not bool(cond.data) for cond in conds):
+        raise op.exc_type(op.msg)
+
     warnings.warn(
-        f"""Skipping `CheckAndRaise` Op (assertion: {op.msg}) as MLX tracing would remove it.""",
+        f"""Skipping `{type(op).__name__}` Op (assertion: {op.msg}) as MLX tracing would remove it.""",
         stacklevel=2,
     )
 

pytensor/link/mlx/dispatch/blockwise.py

@@ -2,105 +2,46 @@
 
 from pytensor.link.mlx.dispatch import mlx_funcify
 from pytensor.tensor.blockwise import Blockwise
-from pytensor.tensor.signal.conv import Convolve1d as Conv1d
-
-
-def blockwise_conv1d(op, node, **kwargs):
-    """
-    Custom implementation of Blockwise.conv1d for MLX.
-    """
-
-    def batched_conv1d(
-        x: mx.array,
-        kernels: mx.array,
-        mode: str = op.core_op.mode,
-        stride: int = 1,
-        dilation: int = 1,
-    ) -> mx.array:
-        """
-        Apply B separate 1D convolutions (full or valid) to B sequences in parallel.
-
-        Parameters
-        ----------
-        x        : array of shape (B, T)
-                B sequences of length T.
-        kernels  : array of shape (B, K)
-                B kernels of length K.
-        mode     : {"valid", "full"}
-                "valid" → no padding, output length = T - K + 1
-                "full"  → zero-pad so output length = T + K - 1
-        stride   : int, convolution stride (default=1)
-        dilation : int, convolution dilation (default=1)
-
-        Returns
-        -------
-        out      : array of shape (B, L)
-                where L =
-                    - T - K + 1   if mode="valid"
-                    - T + K - 1   if mode="full"
-        """
-        # --- 1) shape checks ---
-        B, T = x.shape
-        Bk, K = kernels.shape
-        if B != Bk:
-            raise ValueError(f"Batch mismatch: x has {B}, kernels has {Bk}")
-
-        # --- 2) flip kernels for convolution ---
-        kernels_flipped = kernels[:, ::-1]  # shape (B, K)
-
-        # --- 3) decide padding ---
-        if mode == "valid":
-            pad = 0
-        elif mode == "full":
-            pad = (K - 1) * dilation
-        else:
-            raise ValueError(f"Unsupported mode {mode!r}: choose 'valid' or 'full'")
-
-        # --- 4) reshape into MLX conv1d form ---
-        #   input: (N=1, H=T, C_in=B)
-        x_in = x.T[None, :, :]
-
-        #   weight: (C_out=B, H_f=K, C_in=1)
-        w = kernels_flipped[:, :, None]
-
-        # --- 5) run grouped conv1d ---
-        y = mx.conv1d(x_in, w, stride=stride, padding=pad, dilation=dilation, groups=B)
-        # y shape: (1, H_out, B)
-
-        # --- 6) return shape (B, H_out) ---
-        return y[0].T
-
-    return batched_conv1d
 
 
 @mlx_funcify.register(Blockwise)
 def funcify_Blockwise(op: Blockwise, node, **kwargs):
-    # 1) If it's a Conv1d Blockwise, use the custom implementation
-    if isinstance(op.core_op, Conv1d):
-        return blockwise_conv1d(op, node, **kwargs)
-
     # 2) Otherwise, get the core python function for this Blockwise
     core_node = op._create_dummy_core_node(node.inputs)
     core_f = mlx_funcify(op.core_op, core_node)
 
     # 3) Determine how many inputs correspond to batch dimensions
     n_batch = op.batch_ndim(node)
 
-    # 4) Build in_axes: map only the first n_batch args, keep the rest static
-    in_axes = tuple(0 if i < n_batch else None for i in range(len(node.inputs)))
+    # 4) Handle case where no vectorization is needed
+    if n_batch == 0:
+
+        def blockwise_fun(*inputs):
+            return core_f(*inputs)
+
+        return blockwise_fun
+
+    # 5) Vectorize using mx.vmap over any batched inputs
+    in_axes = []
+    for inp, sig in zip(node.inputs, op.inputs_sig):
+        batch_ndim = inp.type.ndim - len(sig)
+        if batch_ndim == 0:
+            in_axes.append(None)
+            continue
+
+        batch_bcast = inp.type.broadcastable[:batch_ndim]
+        # If all batch dims are broadcastable (size 1), treat input as static
+        in_axes.append(0 if not all(batch_bcast) else None)
+
+    if not any(axis == 0 for axis in in_axes):
 
-    # 5) Handle case where no vectorization is needed
-    if n_batch == 0 or all(axis is None for axis in in_axes):
-        # No batch dimensions, just return the core function
         def blockwise_fun(*inputs):
             return core_f(*inputs)
 
         return blockwise_fun
 
-    # 6) Vectorize (vmap) with in_axes
-    blockwise_f = mx.vmap(core_f, in_axes=in_axes)
+    blockwise_f = mx.vmap(core_f, in_axes=tuple(in_axes))
 
-    # 7) Return the mapped function
     def blockwise_fun(*inputs):
         return blockwise_f(*inputs)
 

pytensor/link/mlx/dispatch/core.py

@@ -32,6 +32,17 @@
 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. "
+    "\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"
+)
+
+
 @mlx_funcify.register(Join)
 def mlx_funcify_Join(op, **kwargs):
     def join(axis, *tensors):
@@ -247,33 +258,66 @@ def allocempty(*shape):
 
 @mlx_funcify.register(Alloc)
 def mlx_funcify_Alloc(op, node, **kwargs):
+    node_inputs = getattr(node, "inputs", None)
+    static_dims = (
+        _extract_static_dims(node_inputs[1:])
+        if node_inputs and len(node_inputs) > 1
+        else None
+    )
+
     def alloc(x, *shape):
-        try:
-            # Convert shape elements to Python ints for MLX compatibility
-            # MLX requires shape dimensions to be Python integers, not MLX arrays
-            shape_ints = tuple(
-                int(s.item()) if hasattr(s, "item") else int(s) for s in shape
-            )
-            return mx.broadcast_to(x, shape_ints)
-        except ValueError as e:
-            if (
-                "[eval] Attempting to eval an array during function transformations"
-                in str(e)
-            ):
-                # This is the MLX compilation limitation - provide helpful error
-                raise ValueError(
-                    "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. "
-                    # Just a note! TODO: remove this once we have a better solution
-                    "\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"
-                    "\n\nOriginal error: " + str(e)
-                ) from e
-            else:
-                # Re-raise other ValueError exceptions
-                raise
+        resolved_shape = (
+            _resolve_shape(static_dims, shape)
+            if static_dims is not None
+            else tuple(_coerce_to_int(dim) for dim in shape)
+        )
+        result = mx.broadcast_to(x, resolved_shape)
+        if node_inputs is not None:
+            value_for_check = x if hasattr(x, "shape") else np.asarray(x)
+            Alloc._check_runtime_broadcast(node, value_for_check, resolved_shape)
+        return result
 
     return alloc
+
+
+def _extract_static_dims(shape_inputs):
+    static_dims = []
+    for dim in shape_inputs:
+        try:
+            static_dims.append(int(get_scalar_constant_value(dim)))
+        except NotScalarConstantError:
+            static_dims.append(None)
+    return tuple(static_dims)
+
+
+def _resolve_shape(static_dims, runtime_shape):
+    if len(static_dims) != len(runtime_shape):
+        raise ValueError("Alloc received unexpected number of shape dimensions")
+
+    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))
+
+    return tuple(resolved)
+
+
+def _coerce_to_int(value):
+    if isinstance(value, np.integer | int):
+        return int(value)
+    try:
+        if hasattr(value, "item"):
+            return int(value.item())
+        return int(value)
+    except (ValueError, TypeError) as exc:
+        _rethrow_dynamic_shape_error(exc)
+        raise
+    raise TypeError(
+        "MLX Alloc expects integer shape components; got value of type "
+        f"{type(value).__name__}."
+    )
+
+
+def _rethrow_dynamic_shape_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

pytensor/link/mlx/dispatch/math.py

@@ -38,7 +38,7 @@
 )
 from pytensor.scalar.math import Erfc, Erfcx, Sigmoid, Softplus
 from pytensor.tensor.elemwise import Elemwise
-from pytensor.tensor.math import Dot
+from pytensor.tensor.math import Argmax, Dot, Max
 
 
 @mlx_funcify.register(Dot)
@@ -49,6 +49,66 @@ def dot(x, y):
     return dot
 
 
+@mlx_funcify.register(Max)
+def mlx_funcify_Max(op, node=None, **kwargs):
+    def max_fn(x):
+        axes = op.axis
+        if axes is None:
+            reduce_axes = None
+        else:
+            reduce_axes = tuple(int(ax) for ax in axes)
+
+        keepdims = getattr(op, "keepdims", False)
+
+        return mx.max(x, axis=reduce_axes, keepdims=keepdims)
+
+    return max_fn
+
+
+@mlx_funcify.register(Argmax)
+def mlx_funcify_Argmax(op, node=None, **kwargs):
+    axis = op.axis
+
+    def argmax_fn(x):
+        if axis is None:
+            axes = tuple(range(x.ndim))
+        else:
+            axes = tuple(int(ax) for ax in axis)
+
+        keep_axes = [i for i in range(x.ndim) if i not in axes]
+        transposed_x = mx.transpose(x, tuple(keep_axes + list(axes)))
+
+        kept_shape = transposed_x.shape[: len(keep_axes)]
+        reduced_shape = transposed_x.shape[len(keep_axes) :]
+
+        flat_size = 1
+        for dim in reduced_shape:
+            flat_size *= int(dim)
+        reshaped_x = transposed_x.reshape((*kept_shape, flat_size))
+
+        max_idx = mx.argmax(reshaped_x, axis=-1)
+
+        result = max_idx.astype(mx.int64)
+
+        if getattr(op, "keepdims", False):
+            reshape_shape = []
+            keep_iter = iter(kept_shape)
+            axis_iter = iter(sorted(axes))
+            next_axis = next(axis_iter, None)
+            for dim_idx in range(x.ndim):
+                if next_axis is not None and dim_idx == next_axis:
+                    reshape_shape.append(1)
+                    next_axis = next(axis_iter, None)
+                else:
+                    reshape_shape.append(int(next(keep_iter)))
+
+            return result.reshape(tuple(reshape_shape))
+
+        return result
+
+    return argmax_fn
+
+
 # Second-level dispatch for scalar operations in Elemwise
 @singledispatch
 def mlx_funcify_Elemwise_scalar_op(scalar_op):