MLX backend diverges from PyTorch in four places (stitching, CPM RevIN refine, residual activation, NaN handling)

I've been working with the MLX backend for TimesFM3 and found four places where it produces different numbers than the PyTorch backend on the same weights and inputs.

To isolate this, I built a tiny model in both backends, copied the PyTorch weights into the MLX model via a safetensors round-trip, and confirmed a default-config decode() call matches (to ~1e-6). Then I flipped one setting at a time:

1. use_stitching=False is parsed but never read. decode() always runs the stitching-based formula:

# mlx/model.py, decode()
extract_len = min(2 * p, cfg.output_patch_len)
overlap = extract_len - p
num_forecast_patches = max(math.ceil((horizon - overlap) / p), 1)
num_hor_patches = num_forecast_patches + cfg.rolls - 1
padded_h = num_hor_patches * p
hor_pad = padded_h - horizon
# cfg.use_stitching is never checked

Max abs diff against PyTorch on my test config: 3.68.

1. use_iterative_cpm_revin doesn't exist in TimesFM3MlxConfig. The refine step runs any time a patch_cpm_mask is present... which is always, on every decode() call:

# mlx/model.py, _forward_logits()
if patch_cpm_mask is not None:
  ref_mu, ref_sigma = cpm_revin_refine_lib.cpm_iterative_revin_refine(
    raw, running_n, mu, sigma, patch_cpm_mask, ...
  )

Max abs diff: 0.053.

1. ResidualBlock hardcodes ReLU, so a swish actually runs ReLU:

# mlx/dense.py
def __call__(self, x: mx.array) -> mx.array:
  return self.output_layer(nn.relu(self.hidden_layer(x))) + self.residual_layer(x)

Max abs diff with activation="swish": 0.997.

1. No NaN sanitization before RevIN stats. PyTorch's forward() does this:

# torch/model.py, forward()
values = torch.nan_to_num(values, nan=0.0)
values = torch.clamp(values, -self.value_clip, self.value_clip)

mlx/model.py's _forward_logits() has no equivalent. It computes running stats straight off the raw input. One unmasked NaN poisons the cumulative mean/std. This is not the case for Pytorch forecasts.

I have fixes for all four, matching PyTorch logic, activation/prenorm/identity_skip support in ResidualBlock, and a NaN sanitize+clip step mirroring torch/model.py) plus a regression test file that checks parity. I'll open a PR with this later this weekend.