Define and implement MLX support for RefitKalmanFilter

[cboulay]

Problem

RefitKalmanFilter and RefitKalmanFilterProcessor describe their core operations as Array API/backend preserving, but the MLX path currently fails before useful work.

Reproduced on an M4 Pro with MLX 0.32.0:

import mlx.core as mx
from ezmsg.learn.model.refit_kalman import RefitKalmanFilter

model = RefitKalmanFilter()
model.fit(mx.random.normal((32, 5)), mx.random.normal((32, 4)))

This fails with:

AttributeError: module 'mlx.core.linalg' has no attribute 'matrix_transpose'

The processor state initialization independently fails because it creates x and P with xp.float64 on the message device:

ValueError: float64 is not supported on the GPU

Additional compatibility and performance concerns

• mlx.core.linalg does not provide matrix_transpose.
• inv and pinv require an explicit MLX CPU stream.
• Checkpoint matrices may remain NumPy while processor state and measurements are MLX, creating a mixed-namespace model.
• The SciPy DARE solve is necessarily a NumPy/CPU boundary and should be treated explicitly.
• Processing loops over samples in Python and mutates filtered_data one row at a time. With MLX this is likely to create a long recurrent lazy graph with many tiny operations.
• State dimensions are normally small, so NumPy/CPU execution may be faster and simpler than GPU dispatch even when messages arrive as MLX arrays.

Affected code is in:

• src/ezmsg/learn/model/refit_kalman.py
• src/ezmsg/learn/process/refit_kalman.py

Design decision needed

Choose and document one supported contract:

1. Backend-preserving MLX state and outputs, with supported dtypes, portable transpose operations, explicit CPU-stream linalg, and bounded recurrent graph materialization; or
2. A deliberate NumPy/CPU internal implementation with one conversion at ingress and, if required by the processor contract, one conversion back to the message namespace at egress.

The second option may be preferable unless profiling shows a benefit from MLX for realistic state sizes and chunk lengths.

Acceptance criteria

• MLX messages no longer fail due to float64, unsupported transpose calls, CPU-only linalg, or mixed namespaces.
• The chosen backend contract is documented.
• Tests cover model fit, checkpoint-loaded processing, predict/update, and refit with MLX-originated data.
• Numerical results match the NumPy path within an appropriate float32 tolerance.
• A benchmark covers realistic state dimensions and multiple chunk lengths, including retained memory across the per-sample recurrence.