I wrote an MLX port of TimesFM 3.0 that runs on Apple Silicon, and checked it against the reference implementation. All 445 checkpoint tensors map one to one, and decode() matches the PyTorch reference to about 1e-6 max abs error from horizon 24 through 512. That range includes the multi-patch path that uses cpm_revin_refine. On an M4 Max the 330M model runs at roughly 20ms per single forecast and about 475 forecasts/sec batched.
Would you be open to an MLX backend living in the repo next to torch/ and flax/, for example under src/timesfm3/mlx/? Amazon did something similar with an mlx branch of chronos-forecasting. I can open a PR under Apache-2.0 with the parity test included.
The current implementation is here (MIT, standalone for now): https://github.com/rachittshah/mlx-tsfm. The parity test is in tests/test_timesfm.py.
I wrote an MLX port of TimesFM 3.0 that runs on Apple Silicon, and checked it against the reference implementation. All 445 checkpoint tensors map one to one, and decode() matches the PyTorch reference to about 1e-6 max abs error from horizon 24 through 512. That range includes the multi-patch path that uses cpm_revin_refine. On an M4 Max the 330M model runs at roughly 20ms per single forecast and about 475 forecasts/sec batched.
Would you be open to an MLX backend living in the repo next to torch/ and flax/, for example under src/timesfm3/mlx/? Amazon did something similar with an mlx branch of chronos-forecasting. I can open a PR under Apache-2.0 with the parity test included.
The current implementation is here (MIT, standalone for now): https://github.com/rachittshah/mlx-tsfm. The parity test is in tests/test_timesfm.py.