From 6ae67d41d813fcdab0a1bc785b79053c3769a63e Mon Sep 17 00:00:00 2001 From: darkpowerxo Date: Thu, 9 Apr 2026 23:00:36 -0400 Subject: [PATCH] revert: drop PR #393 (xreg batch behavior) and PR #390 (SKILL.md link) per maintainer feedback --- .gitignore | 1 + README.md | 5 +- src/timesfm/utils/xreg_lib.py | 110 ++++++++++++++++------------------ 3 files changed, 54 insertions(+), 62 deletions(-) diff --git a/.gitignore b/.gitignore index 24495b0..00dc636 100644 --- a/.gitignore +++ b/.gitignore @@ -8,3 +8,4 @@ datasets/ results/ uv.lock development_setup.md +debug.log diff --git a/README.md b/README.md index 56e8253..c7b2d50 100644 --- a/README.md +++ b/README.md @@ -26,12 +26,11 @@ This open version is not an officially supported Google product. Added fine-tuning example using HuggingFace Transformers + PEFT (LoRA) — see [`timesfm-forecasting/examples/finetuning/`](timesfm-forecasting/examples/finetuning/). -Also added unit tests (`tests/`), fixed per-input ridge regression in XReg to -prevent data leakage, and incorporated several community fixes. +Also added unit tests (`tests/`) and incorporated several community fixes. ## Update - Mar. 19, 2026 -Huge shoutout to [@borealBytes](https://github.com/borealBytes) for adding the support for [AGENTS](https://github.com/google-research/timesfm/blob/master/AGENTS.md)! TimesFM [SKILL.md](https://github.com/google-research/timesfm/blob/master/timesfm-forecasting/SKILL.md) is out. +Huge shoutout to [@borealBytes](https://github.com/borealBytes) for adding the support for [AGENTS](https://github.com/google-research/timesfm/blob/master/AGENTS.md)! TimesFM [SKILL.md](https://github.com/google-research/timesfm/tree/master/timesfm-forecasting) is out. ## Update - Oct. 29, 2025 diff --git a/src/timesfm/utils/xreg_lib.py b/src/timesfm/utils/xreg_lib.py index 2759b67..7a1b19b 100644 --- a/src/timesfm/utils/xreg_lib.py +++ b/src/timesfm/utils/xreg_lib.py @@ -370,20 +370,11 @@ class BatchedInContextXRegBase: x_train = np.concatenate(x_train, axis=1) x_test = np.concatenate(x_test, axis=1) - # Normalize per-input for robustness (batch-wide normalization - # would make each input's result depend on batch composition). - train_splits = np.cumsum(self.train_lens)[:-1] - test_splits = np.cumsum(self.test_lens)[:-1] - train_parts = np.split(x_train, train_splits, axis=0) - test_parts = np.split(x_test, test_splits, axis=0) - norm_train, norm_test = [], [] - for tr, te in zip(train_parts, test_parts): - m = np.mean(tr, axis=0, keepdims=True) - s = np.where((w := np.std(tr, axis=0, keepdims=True)) > _TOL, w, 1.0) - norm_train.append((tr - m) / s) - norm_test.append((te - m) / s) - x_train = [np.concatenate(norm_train, axis=0)] - x_test = [np.concatenate(norm_test, axis=0)] + # Normalize for robustness. + x_mean = np.mean(x_train, axis=0, keepdims=True) + x_std = np.where((w := np.std(x_train, axis=0, keepdims=True)) > _TOL, w, 1.0) + x_train = [(x_train - x_mean) / x_std] + x_test = [(x_test - x_mean) / x_std] # Categorical features. Encode one by one. one_hot_encoder = preprocessing.OneHotEncoder( @@ -472,57 +463,58 @@ class BatchedInContextXRegLinear(BatchedInContextXRegBase): assert_covariate_shapes=assert_covariate_shapes, ) + x_train = x_train_raw.copy() + if max_rows_per_col: + nrows, ncols = x_train.shape + if nrows > (w := ncols * max_rows_per_col): + subsample = jax.random.choice( + jax.random.PRNGKey(max_rows_per_col_sample_seed), + nrows, + (w,), + replace=False, + ) + x_train = x_train[subsample] + flat_targets = flat_targets[subsample] + device = jax.devices("cpu")[0] if force_on_cpu else None + # Runs jitted version of the solvers which are quicker at the cost of + # running jitting during the first time calling. Re-jitting happens whenever + # new (padded) shapes are encountered. + # Ocassionally it helps with the speed and the accuracy if we force single + # thread execution on cpu for accelerator machines: + # 1. Avoid moving data to accelarator memory. + # 2. Avoid precision loss if any. + with jax.default_device(device): + x_train_raw = _to_padded_jax_array(x_train_raw) + x_train = _to_padded_jax_array(x_train) + flat_targets = _to_padded_jax_array(flat_targets) + x_test = _to_padded_jax_array(x_test) + beta_hat = ( + jnp.linalg.pinv( + x_train.T @ x_train + ridge * jnp.eye(x_train.shape[1]), + hermitian=True, + ) + @ x_train.T + @ flat_targets + ) + y_hat = x_test @ beta_hat + y_hat_context = x_train_raw @ beta_hat if debug_info else None + outputs = [] outputs_context = [] - train_idx, test_idx = 0, 0 - with jax.default_device(device): - for trl, tel in zip(self.train_lens, self.test_lens): - x_tr = x_train_raw[train_idx : train_idx + trl] - x_te = x_test[test_idx : test_idx + tel] - y_tr = flat_targets[train_idx : train_idx + trl] - - x_tr_fit = x_tr.copy() - if max_rows_per_col: - nrows, ncols = x_tr_fit.shape - if nrows > (w := ncols * max_rows_per_col): - subsample = jax.random.choice( - jax.random.PRNGKey(max_rows_per_col_sample_seed), - nrows, - (w,), - replace=False, - ) - x_tr_fit = x_tr_fit[subsample] - y_tr = y_tr[subsample] - - x_tr_raw_j = _to_padded_jax_array(x_tr) - x_tr_j = _to_padded_jax_array(x_tr_fit) - y_tr_j = _to_padded_jax_array(y_tr) - x_te_j = _to_padded_jax_array(x_te) - - beta_hat = ( - jnp.linalg.pinv( - x_tr_j.T @ x_tr_j + ridge * jnp.eye(x_tr_j.shape[1]), - hermitian=True, - ) - @ x_tr_j.T - @ y_tr_j + # Reconstruct the ragged 2-dim batched forecasts from flattened linear fits. + train_index, test_index = 0, 0 + for train_index_delta, test_index_delta in zip(self.train_lens, self.test_lens): + outputs.append(np.array(y_hat[test_index : (test_index + test_index_delta)])) + if debug_info: + outputs_context.append( + np.array(y_hat_context[train_index : (train_index + train_index_delta)]) ) - outputs.append(np.array(x_te_j @ beta_hat)[:tel]) - if debug_info: - outputs_context.append(np.array(x_tr_raw_j @ beta_hat)[:trl]) - - train_idx += trl - test_idx += tel + train_index += train_index_delta + test_index += test_index_delta if debug_info: - return ( - outputs, - outputs_context, - _to_padded_jax_array(flat_targets), - _to_padded_jax_array(x_train_raw), - _to_padded_jax_array(x_test), - ) + return outputs, outputs_context, flat_targets, x_train, x_test else: return outputs