MLX backend: every reshape raises TypeError (shape passed as mx.array)

[jessegrabowski]

mlx_funcify_Reshape forwards the shape input straight to mx.reshape, but the linker typifies every input to mx.array and mx.reshape wants a Python sequence of ints, so no reshape works on this backend at all.

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

x = pt.matrix("x", shape=(6, 4), dtype="float32")
xv = np.zeros((6, 4), dtype="float32")

print(pytensor.function([x], x.reshape((24,)), mode="CVM")(xv).shape)  # (24,)
print(pytensor.function([x], x.reshape((24,)), mode="MLX")(xv).shape)
# TypeError: reshape(): incompatible function arguments
#   Invoked with types: mlx.core.array, mlx.core.array

Potential fix (requires testing):

def reshape(x, shp):
    return mx.reshape(x, tuple(shp.tolist()) if isinstance(shp, mx.array) else tuple(shp))

[jessegrabowski]

Confirmed on mlx 0.32.2 with pytensor 3.3.2. This is MLX-version dependent: mx.reshape(x, (mx.array(2), mx.array(3))) is accepted on 0.32.0 and rejected on 0.32.2, so the failure appears and disappears with the MLX pin.

[ricardoV94]

maybe the switch to join and split dims would help