PERF: TensorType.filter is O(N) for foreign array types (mx.array, jax.Array, torch.Tensor)

Part of #2085.

Current implementation

pytensor/tensor/type.py::TensorType.filter only fast-paths np.ndarray and np.memmap. Anything else (including mlx.core.array, jax.Array, torch.Tensor) falls into the trailing else branch:

def filter(self, data, strict=False, allow_downcast=None) -> np.ndarray:
    ...
    elif isinstance(data, np.ndarray) and (data.dtype == self.numpy_dtype):
        ...
    elif strict:
        ...
    else:
        if allow_downcast:
            data = np.asarray(data).astype(self.dtype)
        else:
            if isinstance(data, np.ndarray):
                ...
            elif (allow_downcast is None and isinstance(data, float | np.floating) ...):
                ...
            else:
                converted_data = np.asarray(data, self.dtype)
                if TensorType.values_eq(
                    np.asarray(data), converted_data, force_same_dtype=False
                ):
                    data = converted_data

The last branch does two np.asarray conversions and an element-wise comparison (values_eq runs a == b then np.all). Cost is O(N) in array size.

Reproducer + metrics

import time, mlx.core as mx, numpy as np
from pytensor.tensor.type import TensorType

T = TensorType("float32", shape=(None, None))
rng = np.random.default_rng(0)

for N in (256, 1024, 2048):
    x_np = rng.standard_normal((N, N)).astype(np.float32)
    x_mx = mx.array(x_np); mx.eval(x_mx); mx.synchronize()

    for label, arr in [("np", x_np), ("mx", x_mx)]:
        for _ in range(5): T.filter(arr)
        t0 = time.perf_counter()
        for _ in range(200): T.filter(arr)
        print(f"{N}x{N} filter({label}): {(time.perf_counter()-t0)/200*1e6:8.1f} us")

Output:

 256x256 filter(np):     0.7 us
 256x256 filter(mx):    21.3 us
1024x1024 filter(np):   0.7 us
1024x1024 filter(mx): 280.7 us
2048x2048 filter(np):   0.8 us
2048x2048 filter(mx): 1354.7 us

cProfile confirms 95 % of the cost is values_eq, not the conversion itself (np.asarray(mx.array) is ~18 µs zero-copy via Apple unified memory / DLPack).

End-to-end impact on a single f_mlx(mx_array, mx_array, mx_array) call for sin(x)*exp(y)+z at 2048²:

→ 4.2× speedup on mx.array input from a one-line addition.

Why __dlpack__ and not __array__

mx.array does not implement __array__. It implements __dlpack__ and the buffer protocol, which is what np.asarray uses internally. JAX jax.Array and PyTorch torch.Tensor also expose __dlpack__. Detecting via hasattr(..., "__dlpack__") or hasattr(..., "__array_interface__") covers all three.

Proposed change

File: pytensor/tensor/type.py, top of TensorType.filter:

def filter(self, data, strict=False, allow_downcast=None) -> np.ndarray:
    if isinstance(data, Variable):
        raise TypeError(...)

    # Fast-path for foreign array types implementing the array protocols
    # (mlx.core.array, jax.Array, torch.Tensor, etc.).
    if not isinstance(data, np.ndarray) and (
        hasattr(data, "__dlpack__") or hasattr(data, "__array_interface__")
    ):
        data = np.asarray(data)

    # ... rest unchanged ...

The conversion is zero-copy on Apple unified memory (M-series) and effectively zero-copy for any device array on the same host. Down-stream the existing np.ndarray fast path takes over.

This benefits MLX, JAX, and PyTorch users equally with no MLX-specific imports.

Tests

Add to tests/tensor/test_type.py:

def test_filter_accepts_dlpack_arrays():
    mx = pytest.importorskip("mlx.core")
    T = TensorType("float32", shape=(None, None))
    x = mx.array(np.zeros((4, 4), dtype=np.float32))
    out = T.filter(x)
    assert isinstance(out, np.ndarray)
    assert out.dtype == np.float32

Start with TensorType.filter in pytensor/tensor/type.py and review the existing conversion and dtype-validation branches. Add the foreign-array coverage in tests/tensor/test_type.py, then run test_filter_accepts_dlpack_arrays and confirm the filter returns a NumPy array without the size-dependent comparison overhead.