PERF: minor MLX dispatch optimizations (Cast, DimShuffle, Argmax)

Part of #2085.

Three small dispatcher optimisations that don't move the needle individually but are low-risk wins. mx.compile absorbs most of these into fused Metal kernels, so the gains show up primarily when use_compile=False or for very short graphs. Filing as one issue because they share the same fix shape.

---

A) Cast resolves dtype on every call

Current

pytensor/link/mlx/dispatch/scalar.py lines 72-98:

@mlx_funcify.register(Cast)
def mlx_funcify_Cast(op, **kwargs):
    scalar_op = getattr(op, "scalar_op", op)

    def cast(x):
        dtype = convert_dtype_to_mlx(scalar_op.o_type.dtype)  # called every call
        try:
            return x.astype(dtype)
        except ValueError as e:
            ...

dtype is recomputed via convert_dtype_to_mlx(...) on every invocation, and a try/except ValueError is set up per-call.

Metric

Direct dispatcher call (no actual cast work):

• current: 0.51 µs/call
• pre-resolved: 0.36 µs/call

Inside mx.compile, both compile away — no measurable difference. The win is in the small-graph / overhead-bound path.

Proposed

@mlx_funcify.register(Cast)
def mlx_funcify_Cast(op, **kwargs):
    scalar_op = getattr(op, "scalar_op", op)
    target_dtype = convert_dtype_to_mlx(scalar_op.o_type.dtype)

    def cast(x):
        return x.astype(target_dtype)

    return cast

The try/except ValueError for "is not supported on the GPU" can be dropped because convert_dtype_to_mlx already auto-casts unsupported dtypes (float64→float32, complex128→complex64) at funcify time and emits the warning then. If the user opted out via auto_cast_unsupported=False, they explicitly want the error.

---

B) DimShuffle always uses transpose+reshape, even for pure expand-dims

Current

pytensor/link/mlx/dispatch/scalar.py lines 20-34:

def dimshuffle(x):
    if isinstance(x, int | float) or ...:
        x = mx.array(x)
    res = mx.transpose(x, op._transposition)
    shape = list(res.shape[: len(op.shuffle)])
    for augm in op.augment:
        shape.insert(augm, 1)
    return mx.reshape(res, shape)

When op._transposition is identity (no axis re-ordering, just adding new axes), the mx.transpose is a no-op but still allocates intermediate state and the manual shape-list construction adds Python overhead.

Metric

(128, 256) array, single-axis expand_dims:

• current: 32.9 µs
• mx.expand_dims(x, axis): 29.3 µs
• inside mx.compile: 32.4 µs vs 28.9 µs

A small win uncompiled, marginal compiled. Mostly a code-clarity improvement.

Proposed

@mlx_funcify.register(DimShuffle)
def mlx_funcify_DimShuffle(op, **kwargs):
    transposition = op._transposition
    augment = list(op.augment)
    shuffle_len = len(op.shuffle)
    is_identity_transpose = list(transposition) == list(range(len(transposition)))

    if is_identity_transpose and shuffle_len == len(transposition):
        sorted_aug = sorted(augment)
        def dimshuffle_expand_only(x):
            if isinstance(x, (int, float)) or (
                isinstance(x, np.number) and not isinstance(x, np.ndarray)
            ):
                x = mx.array(x)
            for ax in sorted_aug:
                x = mx.expand_dims(x, ax)
            return x
        return dimshuffle_expand_only

    # General path: keep current logic
    ...

---

C) Argmax does explicit transpose+reshape+argmax(-1) even for single-axis

Current

pytensor/link/mlx/dispatch/math.py lines 31-72: builds keep_axes, mx.transpose, flattens via reshape, calls mx.argmax(axis=-1), reshapes back if keepdims.

Metric

Argmax on (128, 1024, 16) along axis=1:

• current (transpose+reshape): 462 µs uncompiled
• mx.argmax(x, axis=1): 360 µs uncompiled
• compiled with mx.compile: 350 µs both — mx.compile fuses the redundant work away

Real-world: under mx.compile (default for MLX backend) the difference disappears. Listed for completeness; if any code path uses use_compile=False, this saves ~22 %.

Proposed

Add a fast path for the single-axis case (the common case):

@mlx_funcify.register(Argmax)
def mlx_funcify_Argmax(op, node=None, **kwargs):
    axis = op.axis
    keepdims = getattr(op, "keepdims", False)

    if axis is not None and len(axis) == 1:
        single_axis = int(axis[0])
        def argmax_fast(x):
            return mx.argmax(x, axis=single_axis, keepdims=keepdims).astype(mx.int64)
        return argmax_fast

    # multi-axis: keep current transpose+reshape logic
    ...

⚠️ Caveat: during this analysis we observed the single-axis fast path occasionally regressing the np-input pytensor MLX argmax timing (461 µs → 783 µs), reproducible across runs. Cause was not isolated — likely an interaction between dispatcher closure shape and mx.compile fusion at small batch sizes. Recommend microbenchmarking this change against both np and mx inputs before merging, and skipping if the regression reproduces.

---

Files

• pytensor/link/mlx/dispatch/scalar.py (Cast + DimShuffle)
• pytensor/link/mlx/dispatch/math.py (Argmax)

Acceptance criteria

• No regressions in tests/link/mlx/ test suite.
• Microbenchmark for each (Cast, DimShuffle expand-dims, Argmax single-axis) shows neutral or positive change in compiled mode.
• If Argmax regression reproduces (caveat above), drop that part and merge only Cast + DimShuffle.

Start by reading the Cast and DimShuffle implementations in pytensor/link/mlx/dispatch/scalar.py and the Argmax implementation in pytensor/link/mlx/dispatch/math.py. Run the tests/link/mlx/ suite, then microbenchmark Cast, DimShuffle expand-dims, and single-axis Argmax in compiled mode with both NumPy and MLX inputs. Done means no test regressions and neutral or improved benchmarks; omit Argmax if its reported regression remains.