MLX backend: sort and argsort always raise TypeError (axis passed as mx.array)

[guillaume-osmo]

Description

mlx_funcify_Sort and mlx_funcify_ArgSort forward the axis input straight to mx.sort / mx.argsort, but the linker typifies every input to mx.array while both MLX functions require a Python int, so nothing on this backend sorts at all — every call raises TypeError. Same root cause as #2386.

All five tests in tests/link/mlx/test_sort.py are already red on main.

import numpy as np
import pytensor
import pytensor.tensor as pt

x = pt.matrix("x", shape=(3, 4), dtype="float32")
xv = np.random.default_rng(0).normal(size=(3, 4)).astype("float32")

print(pytensor.function([x], pt.sort(x, axis=-1), mode="CVM")(xv).shape)  # (3, 4)
print(pytensor.function([x], pt.sort(x, axis=-1), mode="MLX")(xv).shape)
# TypeError: sort(): incompatible function arguments
#   Invoked with types: mlx.core.array, kwargs = { axis: mlx.core.array }

axis=None fails one step earlier, in Reshape (#2386), since it flattens first.

Potential fix:

def sort(x, axis):
    return mx.sort(x, axis=int(axis))

Found while working on #2385: the gradient of a batched solve reaches argsort through the pivot-to-permutation step, so this also blocks that.