Preserve the signed-log derivative at zero in PyTorch

[sylvesterkaczmarek]

Summary

Fixes #535.

Choose a nonzero branch sign at zero when evaluating the signed-log transform and its inverse. The previous sign(x)*f(abs(x)) expression returned a zero automatic derivative at zero, even though both functions have slope one there.

Preserve forward values, nonzero derivatives, shapes, dtypes, devices, registry lookup and the other transforms. The backward-compatible import path inherits the correction. This is limited to the PyTorch optional signed-log transform; the default identity transform and MLX backend are unchanged.

Validation

Five methods with sixteen subtests extend the existing primitives suite. Tests check analytical derivatives, finite-difference gradcheck, tangent roundtrips, noncontiguous tensors, empty/nonfinite inputs, float16/bfloat16/float32/float64 values and an SGD update from zero initialization. Backward execution is also tested through torch.compile with the aot_eager backend and fullgraph=True.

• New checks reproduce 9 reported regression failures on unchanged source; value controls pass.
• Baseline selected suite: 35 passed.
• python -m pytest src/timesfm3/torch/primitives_test.py src/timesfm3/torch/transformer_test.py src/timesfm3/torch/cpm_revin_refine_test.py src/timesfm3/torch/model_test.py -q: 40 passed, 16 subtests passed.
• Passed on macOS CPU with Python 3.11.16 / PyTorch 2.14.0 and Python 3.12.11 / PyTorch 2.14.1.
• Ruff E4/E7/E9/F, changed-line Ruff formatting, compilation and git diff --check passed.

The model tests use their existing small randomly initialized fixtures. Pretrained forecasting benchmarks, GPU/MPS, TorchInductor and MLX were not tested. No dependencies, workflows or existing test expectations change. This corrects the first derivative; it does not assert higher-order differentiability at the origin.

Based on master at e51928e27119cb17bebc005be2696b75e0a9e688.

[hossiendehghan989]

The fix is well scoped, and the validation matrix is unusually strong for this class of autograd edge case. In particular, testing both analytical and finite-difference gradients, non-contiguous tensors, mixed dtypes, and the zero-initialized SGD step covers the failure modes that a value-only roundtrip test would miss.

One implementation detail worth preserving in the final code is that the zero branch must be explicit rather than relying on torch.sign(0), because that is the source of the zero-gradient regression. I would also keep the forward-value assertions separate from the derivative assertions so a future refactor cannot make the transform numerically correct while reintroducing the tangent failure. The stated limitation around higher-order differentiability is clear and appropriately scoped.

Given the clean merge state and the reported 40 passing tests, this looks ready from the numerical/API side. I would be happy to review any follow-up covering the same invariant in another backend if that becomes necessary.

[sylvesterkaczmarek]

Thanks for checking this. The explicit zero branch and separate value/gradient checks are retained. Could a maintainer approve the pending build and review the PR?