Generates a Tensor of sinusoids with different frequencies. Args: seq_length: an optional Python int defining the output sequence length.

Loads an adapter checkpoint and merges it with the original model weights. Args: model (TimesFm): The model to update. adapt

Loads a PyTorch TimesFM model from a checkpoint.

Loads a TimesFM model from a checkpoint.

Loads a checkpoint and compiles the decoder.

Log metrics to the specified backend. Args: metrics: Dictionary containing metric names and values. step: Optional step n

Main function that selects and runs the appropriate training mode.

Calculates the moving average using NumPy's convolution function.

Scans along an axis.

Construct the model.

Setup process function with optimized CUDA handling.

A small residual block with SiLU/Swish activation (matches TimesFM).

The layer must handle (batch, seq, features) tensors.

All three activation modes must be constructable without error.

All supported activations must produce finite, non-NaN output.

When every element is masked, the function must return zeros rather than NaN or raise an error. This happens when an input series is entirely

When every element is NaN, ``np.argmax`` on an all-False mask returns 0 — so the implementation returns the original array, not an empty one.

Each sample in the batch must be computed independently. Cross-sample leakage would corrupt multi-series forecasting.

A constant series has zero variance — sigma must be exactly 0. This is important because ``revin`` guards against division-by-zero using ``_T

Default config must be conservative: no normalization, no fancy heads. These defaults are what users get when they call ``ForecastConfig()``

Two configs with identical fields must be equal (value semantics).

Leading NaN values are filled with the first valid value. In practice ``strip_leading_nans`` runs first, but the function must still be robus

Trailing NaN values are filled via ``np.interp`` which holds the last known value (nearest-neighbor extrapolation).

Verifies that Flax model wrapper constructor accepts arbitrary kwargs.

After forward normalization: (x - mu) / sigma.

normalize → denormalize must reconstruct the original tensor. This is the fundamental invariant of reversible normalization: the model operat

cos and sin outputs are bounded in [-1, 1]; sign outputs are bounded in {-1, 0, 1}. The total Fourier part (before residual) is thus bounded.

Configs are frozen dataclasses — mutating them must raise. This is critical because ``compile()`` captures the config object and the compiled

Gradients must reach both the main path and the residual path. Dead gradients on either path would prevent the layer from learning.

Accumulating two batches incrementally must yield the same result as computing statistics over all values at once. This is the defining prope

A single differing field must break equality.

Multiple consecutive interior NaN values are interpolated.

A single interior NaN is linearly interpolated from neighbors.

If the known values are strictly increasing, the interpolated result must also be non-decreasing — a basic sanity check on the interpolation d

Unsupported activation must raise ``ValueError`` immediately — fail fast rather than producing garbage at inference time.

DecodeCache is *not* frozen — the attention loop mutates ``next_index`` and ``num_masked`` in-place during autoregressive decoding.

Key and value tensors must have identical shapes — they are indexed in parallel during attention computation.

Masked positions must be completely ignored — as if they don't exist. In TimesFM, leading padding is masked. If mask handling is broken, the

The model instantiation will fail if this invariant is broken. We verify the config at least *carries* the right values that the TimesFM 2.5

Sigma values just below ``_TOLERANCE`` must trigger the guard.

RevIN must work for series with negative values (e.g. temperature, financial returns). ``infer_is_positive`` is a separate downstream flag and

If the first element is valid, nothing is stripped regardless of internal NaNs.

A zero-valued input must not cause NaN (epsilon prevents div-by-0).

Without NaN values the array is returned as-is (fast path).

An array without NaN values must pass through unmodified.

The four Fourier components (cos, sin, sq_wave_1, sq_wave_2) require ``output_dims`` to be divisible by 4.

After interpolation, no NaN values should remain.

Output must have the config's ``output_dims`` as the last dimension, regardless of input batch shape.

Output dims must be exactly ``config.output_dims``.

The block must handle (batch, seq, features) inputs — the layout used when processing patched time series.

RMSNorm must not change the tensor shape.

Output dtype must match input dtype (float32 stays float32).

Non-NaN values in the original array must never be modified.

``dataclasses.replace`` must yield a new object with updated fields. The compile path uses ``replace`` to adjust context/horizon to valid mul

The residual connection must contribute to the output. We verify this by comparing the output when the hidden path is zeroed out vs. the full

After reverse: x * sigma + mu.

RevIN must broadcast correctly for (batch, patches, patch_len) tensors — the actual shape used during patched decoding.

RevIN must broadcast correctly for (batch, patches, patch_len, q) tensors — the shape used for quantile outputs. In the actual decode path, m

A single update with no mask must match numpy's mean and std. This is the most basic correctness check: feed all values at once and compare a

Edge case: exactly one leading NaN.

With only one valid value, every NaN is replaced by that value (np.interp clamps to the single known point).

A single non-NaN element must be returned as-is.

StackedTransformersConfig must wrap a TransformerConfig cleanly.

Leading NaNs are removed; NaNs embedded in the middle are kept.

Verifies that torch_compile=True compiles model.forward, not a no-op.

Verifies that PyTorch load_checkpoint and from_pretrained work locally.

Verifies that torch_compile=False leaves model.forward as a plain method.

With scale = 1, output should have approximately unit RMS along the feature dimension — that's the point of RMS normalization.

With default scale (initialized to zeros), output must be all zeros. This is a critical initialization property: at init, each transformer la

When sigma < tolerance, the function substitutes 1.0 to avoid division by zero. This occurs for constant-valued input series. NaN propagation