Reversible instance normalization.

Reversible per-instance normalization.

Performs linear interpolation to fill NaN values in a 1D numpy array. Args: arr: The 1D numpy array containing NaN values. Returns:

Removes contiguous NaN values from the beginning of a NumPy array. Args: arr: The input NumPy array. Returns: A new NumPy array with lea

Updates the running stats.

Forecasts on a list of time series. Args: inputs: list of time series forecast contexts. Each context time series should be in a fo

Loads a Flax TimesFM model.

Finish the W&B run if on the main process.

Return total physical RAM in GB, cross-platform.

Loads a TimesFM model from a checkpoint directory or file.

Return available RAM in GB.

Saves the model's state dictionary to a safetensors file. This method is called by the `save_pretrained` method from `PyTorchModelHubMixin`.

Decodes the time series.

Returns a large negative value for the given dtype.

Loads a checkpoint and compiles the decoder.

Tensorflow Dataset.

Create appropriate DataLoader based on training configuration. Args: dataset: Dataset to create loader for. is_train: Whe

Get features and ts in specified windows.

Formats and pads raw inputs to feed into the model. This function both pads each time series to match the context length, and pads the inputs

Process a single batch of data. Args: batch: List of input tensors. Returns: Tuple of (loss, predictions).

Scans along an axis.

Check if a Python package is installed.

Compiles the TimesFM model for fast decoding.

Estimate memory requirements for a dataset. Args: num_series: Number of time series in the dataset context_length: Length of each

Compare zero-shot vs fine-tuned on a subset of stores.

Train the model. Args: train_dataset: Training dataset. val_dataset: Validation dataset. Returns: Dict

Forecasts the time series.

Forecasts on a list of time series. Args: inputs: A pd.DataFrame of all time series. The dataframe should have a `unique_id` column

Jitting decoding function.

Updates target modules with adapter layers. Args: mdl_vars (dict): Model variables. model (pax_fiddle.Config): Model configu

Makes attention mask.

Makes attention mask.

Plot model predictions against ground truth for a batch of validation data. Args: model: Trained TimesFM model val_dataset: Vali

Print a detailed memory estimate for a dataset. Args: num_series: Number of time series in the dataset context_length: Length of

Run all system checks and return a report.

Additional initialization for subclasses before checkpoint loading.

All Jax steps after model decode call.

Verifies the validity of the covariate inputs.

Verifies the validity of the covariate inputs.

All Jax steps before model decode call.

Return distance to given holiday.

Fixes quantile crossing.

Forces flip invariance.

Forecasts on a list of time series. Args: inputs: list of time series forecast contexts. Each context time series should be in a fo

Input is of shape [B, N, P].

Input is of shape [B, N, P].

Extracts adapter weight parameters from the given variable weight hyperparameters. Args: var_weight_hparams (dict): Variable weight

Get categorical columns.

Initializes and adds adapter parameters to target modules. Args: mdl_vars (dict): Model variables. num_layers (int): Number

mae loss.