tsagentkit-timesfm 1.0.0__py3-none-any.whl

timesfm/utils/xreg_lib.py

@@ -0,0 +1,520 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Helper functions for in-context covariates and regression."""
+
+import itertools
+import math
+from typing import Any, Iterable, Literal, Mapping, Sequence
+
+try:
+  import jax
+  import jax.numpy as jnp
+  import numpy as np
+  from sklearn import preprocessing
+except ImportError:
+  raise ImportError(
+    "Failed to load the XReg module. Did you forget to install `timesfm[xreg]`?"
+  )
+
+Category = int | str
+
+_TOL = 1e-6
+XRegMode = Literal["timesfm + xreg", "xreg + timesfm"]
+
+
+def _unnest(nested: Sequence[Sequence[Any]]) -> np.ndarray:
+  return np.array(list(itertools.chain.from_iterable(nested)))
+
+
+def _repeat(elements: Iterable[Any], counts: Iterable[int]) -> np.ndarray:
+  return np.array(
+    list(itertools.chain.from_iterable(map(itertools.repeat, elements, counts)))
+  )
+
+
+def _to_padded_jax_array(x: np.ndarray) -> jax.Array:
+  if x.ndim == 1:
+    (i,) = x.shape
+    di = 2 ** math.ceil(math.log2(i)) - i
+    return jnp.pad(x, ((0, di),), mode="constant", constant_values=0.0)
+  elif x.ndim == 2:
+    i, j = x.shape
+    di = 2 ** math.ceil(math.log2(i)) - i
+    dj = 2 ** math.ceil(math.log2(j)) - j
+    return jnp.pad(x, ((0, di), (0, dj)), mode="constant", constant_values=0.0)
+  else:
+    raise ValueError(f"Unsupported array shape: {x.shape}")
+
+
+# Per time series normalization: forward.
+def normalize(batch):
+  stats = [(np.mean(x), np.where((w := np.std(x)) > _TOL, w, 1.0)) for x in batch]
+  new_batch = [(x - stat[0]) / stat[1] for x, stat in zip(batch, stats)]
+  return new_batch, stats
+
+
+# Per time series normalization: inverse.
+def renormalize(batch, stats):
+  return [x * stat[1] + stat[0] for x, stat in zip(batch, stats)]
+
+
+class BatchedInContextXRegBase:
+  """Helper class for in-context regression covariate formatting.
+
+  Attributes:
+    targets: List of targets (responses) of the in-context regression.
+    train_lens: List of lengths of each target vector from the context.
+    test_lens: List of lengths of each forecast horizon.
+    train_dynamic_numerical_covariates: Dict of covariate names mapping to the
+      dynamic numerical covariates of each forecast task on the context. Their
+      lengths should match the corresponding lengths in `train_lens`.
+    train_dynamic_categorical_covariates: Dict of covariate names mapping to the
+      dynamic categorical covariates of each forecast task on the context. Their
+      lengths should match the corresponding lengths in `train_lens`.
+    test_dynamic_numerical_covariates: Dict of covariate names mapping to the
+      dynamic numerical covariates of each forecast task on the horizon. Their
+      lengths should match the corresponding lengths in `test_lens`.
+    test_dynamic_categorical_covariates: Dict of covariate names mapping to the
+      dynamic categorical covariates of each forecast task on the horizon. Their
+      lengths should match the corresponding lengths in `test_lens`.
+    static_numerical_covariates: Dict of covariate names mapping to the static
+      numerical covariates of each forecast task.
+    static_categorical_covariates: Dict of covariate names mapping to the static
+      categorical covariates of each forecast task.
+  """
+
+  def __init__(
+    self,
+    targets: Sequence[Sequence[float]],
+    train_lens: Sequence[int],
+    test_lens: Sequence[int],
+    train_dynamic_numerical_covariates: (
+      Mapping[str, Sequence[Sequence[float]]] | None
+    ) = None,
+    train_dynamic_categorical_covariates: (
+      Mapping[str, Sequence[Sequence[Category]]] | None
+    ) = None,
+    test_dynamic_numerical_covariates: (
+      Mapping[str, Sequence[Sequence[float]]] | None
+    ) = None,
+    test_dynamic_categorical_covariates: (
+      Mapping[str, Sequence[Sequence[Category]]] | None
+    ) = None,
+    static_numerical_covariates: Mapping[str, Sequence[float]] | None = None,
+    static_categorical_covariates: (Mapping[str, Sequence[Category]] | None) = None,
+  ) -> None:
+    """Initializes with the exogenous covariate inputs.
+
+    Here we use model fitting language to refer to the context as 'train' and
+    the horizon as 'test'. We assume batched inputs. To properly format the
+    request:
+
+     - `train_lens` represents the contexts in the batch. Targets and all train
+     dynamic covariates should have the same lengths as the corresponding
+     elements
+     in `train_lens`. Notice each `train_len` can be different from the exact
+     length of the corresponding context depending on how much of the context is
+     used for fitting the in-context model.
+     - `test_lens` represents the horizon lengths in the batch. All tesdt
+     dynamic
+     covariates should have the same lengths as the corresponding elements in
+     `test_lens`.
+     - Static covariates should be one for each input.
+     - For train and test dynamic covariates, they should have the same
+     covariate
+     names.
+
+     Pass an empty dict {} for a covariate type if it is not present.
+
+     Example:
+       Here is a set of valid inputs whose schema can be used for reference.
+       ```
+       targets = [
+           [0.0, 0.1, 0.2],
+           [0.0, 0.1, 0.2, 0.3],
+       ]  # Two inputs in this batch.
+       train_lens = [3, 4]
+       test_lens = [2, 5]  # Forecast horizons 2 and 5 respectively.
+       train_dynamic_numerical_covariates = {
+           "cov_1_dn": [[0.0, 0.5, 1.0], [0.0, 0.5, 1.0, 1.5]],
+           "cov_2_dn": [[0.0, 1.5, 1.0], [0.0, 1.5, 1.0, 2.5]],
+       }  # Each train dynamic covariate has 3 and 4 elements respectively.
+       test_dynamic_numerical_covariates = {
+           "cov_1_dn": [[0.1, 0.6], [0.1, 0.6, 1.1, 1.6, 2.4]],
+           "cov_2_dn": [[0.1, 1.1], [0.1, 1.6, 1.1, 2.6, 10.0]],
+       }  # Each test dynamic covariate has 2 and 5 elements respectively.
+       train_dynamic_categorical_covariates = {
+           "cov_1_dc": [[0, 1, 0], [0, 1, 2, 3]],
+           "cov_2_dc": [["good", "bad", "good"], ["good", "good", "bad",
+           "bad"]],
+       }
+       test_dynamic_categorical_covariates = {
+           "cov_1_dc": [[1, 0], [1, 0, 2, 3, 1]],
+           "cov_2_dc": [["bad", "good"], ["bad", "bad", "bad", "bad", "bad"]],
+       }
+       static_numerical_covariates = {
+           "cov_1_sn": [0.0, 3.0],
+           "cov_2_sn": [2.0, 1.0],
+           "cov_3_sn": [1.0, 2.0],
+       }  # Each static covariate has 1 element for each input.
+       static_categorical_covariates = {
+           "cov_1_sc": ["apple", "orange"],
+           "cov_2_sc": [2, 3],
+       }
+       ```
+
+    Args:
+      targets: List of targets (responses) of the in-context regression.
+      train_lens: List of lengths of each target vector from the context.
+      test_lens: List of lengths of each forecast horizon.
+      train_dynamic_numerical_covariates: Dict of covariate names mapping to the
+        dynamic numerical covariates of each forecast task on the context. Their
+        lengths should match the corresponding lengths in `train_lens`.
+      train_dynamic_categorical_covariates: Dict of covariate names mapping to
+        the dynamic categorical covariates of each forecast task on the context.
+        Their lengths should match the corresponding lengths in `train_lens`.
+      test_dynamic_numerical_covariates: Dict of covariate names mapping to the
+        dynamic numerical covariates of each forecast task on the horizon. Their
+        lengths should match the corresponding lengths in `test_lens`.
+      test_dynamic_categorical_covariates: Dict of covariate names mapping to
+        the dynamic categorical covariates of each forecast task on the horizon.
+        Their lengths should match the corresponding lengths in `test_lens`.
+      static_numerical_covariates: Dict of covariate names mapping to the static
+        numerical covariates of each forecast task.
+      static_categorical_covariates: Dict of covariate names mapping to the
+        static categorical covariates of each forecast task.
+    """
+    self.targets = targets
+    self.train_lens = train_lens
+    self.test_lens = test_lens
+    self.train_dynamic_numerical_covariates = train_dynamic_numerical_covariates or {}
+    self.train_dynamic_categorical_covariates = (
+      train_dynamic_categorical_covariates or {}
+    )
+    self.test_dynamic_numerical_covariates = test_dynamic_numerical_covariates or {}
+    self.test_dynamic_categorical_covariates = test_dynamic_categorical_covariates or {}
+    self.static_numerical_covariates = static_numerical_covariates or {}
+    self.static_categorical_covariates = static_categorical_covariates or {}
+
+  def _assert_covariates(self, assert_covariate_shapes: bool = False) -> None:
+    """Verifies the validity of the covariate inputs."""
+
+    # Check presence.
+    if (
+      self.train_dynamic_numerical_covariates
+      and not self.test_dynamic_numerical_covariates
+    ) or (
+      not self.train_dynamic_numerical_covariates
+      and self.test_dynamic_numerical_covariates
+    ):
+      raise ValueError(
+        "train_dynamic_numerical_covariates and"
+        " test_dynamic_numerical_covariates must be both present or both"
+        " absent."
+      )
+
+    if (
+      self.train_dynamic_categorical_covariates
+      and not self.test_dynamic_categorical_covariates
+    ) or (
+      not self.train_dynamic_categorical_covariates
+      and self.test_dynamic_categorical_covariates
+    ):
+      raise ValueError(
+        "train_dynamic_categorical_covariates and"
+        " test_dynamic_categorical_covariates must be both present or both"
+        " absent."
+      )
+
+    # Check keys.
+    for dict_a, dict_b, dict_a_name, dict_b_name in (
+      (
+        self.train_dynamic_numerical_covariates,
+        self.test_dynamic_numerical_covariates,
+        "train_dynamic_numerical_covariates",
+        "test_dynamic_numerical_covariates",
+      ),
+      (
+        self.train_dynamic_categorical_covariates,
+        self.test_dynamic_categorical_covariates,
+        "train_dynamic_categorical_covariates",
+        "test_dynamic_categorical_covariates",
+      ),
+    ):
+      if w := set(dict_a.keys()) - set(dict_b.keys()):
+        raise ValueError(f"{dict_a_name} has keys not present in {dict_b_name}: {w}")
+      if w := set(dict_b.keys()) - set(dict_a.keys()):
+        raise ValueError(f"{dict_b_name} has keys not present in {dict_a_name}: {w}")
+
+    # Check shapes.
+    if assert_covariate_shapes:
+      if len(self.targets) != len(self.train_lens):
+        raise ValueError(
+          "targets and train_lens must have the same number of elements."
+        )
+
+      if len(self.train_lens) != len(self.test_lens):
+        raise ValueError(
+          "train_lens and test_lens must have the same number of elements."
+        )
+
+      for i, (target, train_len) in enumerate(zip(self.targets, self.train_lens)):
+        if len(target) != train_len:
+          raise ValueError(
+            f"targets[{i}] has length {len(target)} != expected {train_len}."
+          )
+
+      for key, values in self.static_numerical_covariates.items():
+        if len(values) != len(self.train_lens):
+          raise ValueError(
+            f"static_numerical_covariates has key {key} with number of"
+            f" examples {len(values)} != expected {len(self.train_lens)}."
+          )
+
+      for key, values in self.static_categorical_covariates.items():
+        if len(values) != len(self.train_lens):
+          raise ValueError(
+            f"static_categorical_covariates has key {key} with number of"
+            f" examples {len(values)} != expected {len(self.train_lens)}."
+          )
+
+      for lens, dict_cov, dict_cov_name in (
+        (
+          self.train_lens,
+          self.train_dynamic_numerical_covariates,
+          "train_dynamic_numerical_covariates",
+        ),
+        (
+          self.train_lens,
+          self.train_dynamic_categorical_covariates,
+          "train_dynamic_categorical_covariates",
+        ),
+        (
+          self.test_lens,
+          self.test_dynamic_numerical_covariates,
+          "test_dynamic_numerical_covariates",
+        ),
+        (
+          self.test_lens,
+          self.test_dynamic_categorical_covariates,
+          "test_dynamic_categorical_covariates",
+        ),
+      ):
+        for key, cov_values in dict_cov.items():
+          if len(cov_values) != len(lens):
+            raise ValueError(
+              f"{dict_cov_name} has key {key} with number of examples"
+              f" {len(cov_values)} != expected {len(lens)}."
+            )
+          for i, cov_value in enumerate(cov_values):
+            if len(cov_value) != lens[i]:
+              raise ValueError(
+                f"{dict_cov_name} has key {key} with its {i}-th example"
+                f" length {len(cov_value)} != expected {lens[i]}."
+              )
+
+  def create_covariate_matrix(
+    self,
+    one_hot_encoder_drop: str | None = "first",
+    use_intercept: bool = True,
+    assert_covariates: bool = False,
+    assert_covariate_shapes: bool = False,
+  ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Creates target vector and covariate matrices for in context regression.
+
+    Here we use model fitting language to refer to the context as 'train' and
+    the horizon as 'test'.
+
+    Args:
+      one_hot_encoder_drop: Which drop strategy to use for the one hot encoder.
+      use_intercept: Whether to prepare an intercept (all 1) column in the
+        matrices.
+      assert_covariates: Whether to assert the validity of the covariate inputs.
+      assert_covariate_shapes: Whether to assert the shapes of the covariate
+        inputs when `assert_covariates` is True.
+
+    Returns:
+      A tuple of the target vector, the covariate matrix for the context, and
+      the covariate matrix for the horizon.
+    """
+    if assert_covariates:
+      self._assert_covariates(assert_covariate_shapes)
+
+    x_train, x_test = [], []
+
+    # Numerical features.
+    for name in sorted(self.train_dynamic_numerical_covariates):
+      x_train.append(
+        _unnest(self.train_dynamic_numerical_covariates[name])[:, np.newaxis]
+      )
+      x_test.append(
+        _unnest(self.test_dynamic_numerical_covariates[name])[:, np.newaxis]
+      )
+
+    for covs in self.static_numerical_covariates.values():
+      x_train.append(_repeat(covs, self.train_lens)[:, np.newaxis])
+      x_test.append(_repeat(covs, self.test_lens)[:, np.newaxis])
+
+    if x_train:
+      x_train = np.concatenate(x_train, axis=1)
+      x_test = np.concatenate(x_test, axis=1)
+
+      # 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(
+      drop=one_hot_encoder_drop,
+      sparse_output=False,
+      handle_unknown="ignore",
+    )
+    for name in sorted(self.train_dynamic_categorical_covariates.keys()):
+      ohe_train = _unnest(self.train_dynamic_categorical_covariates[name])[
+        :, np.newaxis
+      ]
+      ohe_test = _unnest(self.test_dynamic_categorical_covariates[name])[:, np.newaxis]
+      x_train.append(np.array(one_hot_encoder.fit_transform(ohe_train)))
+      x_test.append(np.array(one_hot_encoder.transform(ohe_test)))
+
+    for covs in self.static_categorical_covariates.values():
+      ohe = one_hot_encoder.fit_transform(np.array(covs)[:, np.newaxis])
+      x_train.append(_repeat(ohe, self.train_lens))
+      x_test.append(_repeat(ohe, self.test_lens))
+
+    x_train = np.concatenate(x_train, axis=1)
+    x_test = np.concatenate(x_test, axis=1)
+
+    if use_intercept:
+      x_train = np.pad(x_train, ((0, 0), (1, 0)), constant_values=1.0)
+      x_test = np.pad(x_test, ((0, 0), (1, 0)), constant_values=1.0)
+
+    return _unnest(self.targets), x_train, x_test
+
+  def fit(self) -> Any:
+    raise NotImplementedError("Fit is not implemented.")
+
+
+class BatchedInContextXRegLinear(BatchedInContextXRegBase):
+  """Linear in-context regression model."""
+
+  def fit(
+    self,
+    ridge: float = 0.0,
+    one_hot_encoder_drop: str | None = "first",
+    use_intercept: bool = True,
+    force_on_cpu: bool = False,
+    max_rows_per_col: int = 0,
+    max_rows_per_col_sample_seed: int = 42,
+    debug_info: bool = False,
+    assert_covariates: bool = False,
+    assert_covariate_shapes: bool = False,
+  ) -> (
+    list[np.ndarray]
+    | tuple[list[np.ndarray], list[np.ndarray], jax.Array, jax.Array, jax.Array]
+  ):
+    """Fits a linear model for in-context regression.
+
+    Args:
+      ridge: A non-negative value for specifying the ridge regression penalty.
+        If 0 is provided, fallback to ordinary least squares. Note this penalty
+        is added to the normalized covariate matrix.
+      one_hot_encoder_drop: Which drop strategy to use for the one hot encoder.
+      use_intercept: Whether to prepare an intercept (all 1) column in the
+        matrices.
+      force_on_cpu: Whether to force execution on cpu for accelerator machines.
+      max_rows_per_col: How many rows to subsample per column. 0 for no
+        subsampling. This is for speeding up model fitting.
+      max_rows_per_col_sample_seed: The seed for the subsampling if needed by
+        `max_rows_per_col`.
+      debug_info: Whether to return debug info.
+      assert_covariates: Whether to assert the validity of the covariate inputs.
+      assert_covariate_shapes: Whether to assert the shapes of the covariate
+        inputs when `assert_covariates` is True.
+
+    Returns:
+      If `debug_info` is False:
+        The linear fits on the horizon.
+      If `debug_info` is True:
+        A tuple of:
+        - the linear fits on the horizon,
+        - the linear fits on the context,
+        - the flattened target vector,
+        - the covariate matrix for the context, and
+        - the covariate matrix for the horizon.
+    """
+    flat_targets, x_train_raw, x_test = self.create_covariate_matrix(
+      one_hot_encoder_drop=one_hot_encoder_drop,
+      use_intercept=use_intercept,
+      assert_covariates=assert_covariates,
+      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 = []
+
+    # 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)])
+        )
+      train_index += train_index_delta
+      test_index += test_index_delta
+
+    if debug_info:
+      return outputs, outputs_context, flat_targets, x_train, x_test
+    else:
+      return outputs

tsagentkit_timesfm-1.0.0.dist-info/METADATA

@@ -0,0 +1,152 @@
+Metadata-Version: 2.4
+Name: tsagentkit-timesfm
+Version: 1.0.0
+Summary: TimesFM fork distribution for tsagentkit.
+Author-email: Rajat Sen <senrajat@google.com>, Yichen Zhou <yichenzhou@google.com>, Abhimanyu Das <abhidas@google.com>, Petros Mol <pmol@google.com>, Michael Chertushkin <chertushkinmichael@gmail.com>
+Maintainer-email: cuizhengliang <czlbou2012@gmail.com>
+License: Apache-2.0
+Project-URL: Homepage, https://github.com/LeonEthan/timesfm
+Project-URL: Repository, https://github.com/LeonEthan/timesfm
+Project-URL: Issues, https://github.com/LeonEthan/timesfm/issues
+Project-URL: Upstream, https://github.com/google-research/timesfm
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=1.26.4
+Requires-Dist: huggingface_hub[cli]>=0.23.0
+Requires-Dist: safetensors>=0.5.3
+Provides-Extra: torch
+Requires-Dist: torch>=2.0.0; extra == "torch"
+Provides-Extra: flax
+Requires-Dist: flax; extra == "flax"
+Requires-Dist: optax; extra == "flax"
+Requires-Dist: einshape; extra == "flax"
+Requires-Dist: orbax-checkpoint; extra == "flax"
+Requires-Dist: jaxtyping; extra == "flax"
+Requires-Dist: jax[cuda]; extra == "flax"
+Provides-Extra: xreg
+Requires-Dist: jax[cuda]; extra == "xreg"
+Requires-Dist: scikit-learn; extra == "xreg"
+Dynamic: license-file
+
+# TimesFM
+
+TimesFM (Time Series Foundation Model) is a pretrained time-series foundation
+model developed by Google Research for time-series forecasting.
+
+*   Paper:
+    [A decoder-only foundation model for time-series forecasting](https://arxiv.org/abs/2310.10688),
+    ICML 2024.
+*   All checkpoints:
+    [TimesFM Hugging Face Collection](https://huggingface.co/collections/google/timesfm-release-66e4be5fdb56e960c1e482a6).
+*   [Google Research blog](https://research.google/blog/a-decoder-only-foundation-model-for-time-series-forecasting/).
+*   [TimesFM in BigQuery](https://cloud.google.com/bigquery/docs/timesfm-model):
+    an official Google product.
+
+This open version is not an officially supported Google product.
+
+## Fork Notice (tsagentkit)
+
+This repository is a fork of `google-research/timesfm`, published to PyPI as
+`tsagentkit-timesfm` to support the `tsagentkit` project without a GitHub
+dependency. The Python import name remains `timesfm`.
+
+**Latest Model Version:** TimesFM 2.5
+
+**Archived Model Versions:**
+
+-   1.0 and 2.0: relevant code archived in the sub directory `v1`. You can `pip
+    install timesfm==1.3.0` to install an older version of this package to load
+    them.
+
+## Update - Oct. 29, 2025
+
+Added back the covariate support through XReg for TimesFM 2.5.
+
+
+## Update - Sept. 15, 2025
+
+TimesFM 2.5 is out!
+
+Comparing to TimesFM 2.0, this new 2.5 model:
+
+-   uses 200M parameters, down from 500M.
+-   supports up to 16k context length, up from 2048.
+-   supports continuous quantile forecast up to 1k horizon via an optional 30M
+    quantile head.
+-   gets rid of the `frequency` indicator.
+-   has a couple of new forecasting flags.
+
+Along with the model upgrade we have also upgraded the inference API. This repo
+will be under construction over the next few weeks to
+
+1.  add support for an upcoming Flax version of the model (faster inference).
+2.  add back covariate support.
+3.  populate more docstrings, docs and notebook.
+
+### Install
+
+1.  Install from PyPI:
+    ```shell
+    pip install "tsagentkit-timesfm[torch]"
+    # Or with flax
+    pip install "tsagentkit-timesfm[flax]"
+    # Or XReg support
+    pip install "tsagentkit-timesfm[xreg]"
+    ```
+
+2. [Optional] Install your preferred `torch` / `jax` backend based on your OS and accelerators
+(CPU, GPU, TPU or Apple Silicon).:
+
+-   [Install PyTorch](https://pytorch.org/get-started/locally/).
+-   [Install Jax](https://docs.jax.dev/en/latest/installation.html#installation)
+    for Flax.
+
+### From Source (optional)
+
+```shell
+git clone https://github.com/LeonEthan/timesfm.git
+cd timesfm
+
+uv venv
+source .venv/bin/activate
+
+uv pip install -e .[torch]
+# Or with flax
+uv pip install -e .[flax]
+# Or XReg support
+uv pip install -e .[xreg]
+```
+
+### Code Example
+
+```python
+import torch
+import numpy as np
+import timesfm
+
+torch.set_float32_matmul_precision("high")
+
+model = timesfm.TimesFM_2p5_200M_torch.from_pretrained("google/timesfm-2.5-200m-pytorch")
+
+model.compile(
+    timesfm.ForecastConfig(
+        max_context=1024,
+        max_horizon=256,
+        normalize_inputs=True,
+        use_continuous_quantile_head=True,
+        force_flip_invariance=True,
+        infer_is_positive=True,
+        fix_quantile_crossing=True,
+    )
+)
+point_forecast, quantile_forecast = model.forecast(
+    horizon=12,
+    inputs=[
+        np.linspace(0, 1, 100),
+        np.sin(np.linspace(0, 20, 67)),
+    ],  # Two dummy inputs
+)
+point_forecast.shape  # (2, 12)
+quantile_forecast.shape  # (2, 12, 10): mean, then 10th to 90th quantiles.
+```

tsagentkit_timesfm-1.0.0.dist-info/RECORD

@@ -0,0 +1,21 @@
+timesfm/__init__.py,sha256=B2pgyjc3hgIwQQqs5BuJh9tZy3cv1ox39bhyp_JIJBs,919
+timesfm/configs.py,sha256=hhAsa7LEZgEct-gZqP6xl7HBKWR5gHNgHguE7JPjJxk,3617
+timesfm/flax/__init__.py,sha256=Q9FlRO2IfSE9yEaiAYzWkOMBJPCaNYqh4ihcp0t0BQs,574
+timesfm/flax/dense.py,sha256=vkcDueNfcXgR1HL61vCOWzwHCKH9TreiQANJzMK2_Pg,3522
+timesfm/flax/normalization.py,sha256=KGeCNAftZVXtE_5_pfBecBNjJzDMp6JOCl2WoPdBvEM,2168
+timesfm/flax/transformer.py,sha256=K9qO31l5bwUf9zl-myl8HDltAkC-O4CoSWg8up7vIm4,11506
+timesfm/flax/util.py,sha256=2SFLO6xHT84whuNumVIwZTXXPag-XKrKK2hbj74grmM,3027
+timesfm/timesfm_2p5/timesfm_2p5_base.py,sha256=hR-BIcNki2DVJ5k4B3FczT4C0bowpcJu381XR9FLwfw,15174
+timesfm/timesfm_2p5/timesfm_2p5_flax.py,sha256=ts42ros5VP-IdldYqWtWqG9ckNm30Z0dEaXuJ_r9UAo,19375
+timesfm/timesfm_2p5/timesfm_2p5_torch.py,sha256=MPwE6LL_6OOjKJ-qSINO4zGQB_Vho7LfUYjDeoAvboI,16943
+timesfm/torch/__init__.py,sha256=Q9FlRO2IfSE9yEaiAYzWkOMBJPCaNYqh4ihcp0t0BQs,574
+timesfm/torch/dense.py,sha256=KIlJhThOqj3aexuILVu4Qm2LJ0stmgHLFMRlHAWE7B0,3109
+timesfm/torch/normalization.py,sha256=w3tWr9GYrr9e9x7gTvJEpjMY2i5KBJfsmzGQiPbMUVQ,1204
+timesfm/torch/transformer.py,sha256=dEoOHkHBJiFQmY-8Yhu35LsvucebD-RucN6nRFEZYrs,11720
+timesfm/torch/util.py,sha256=xwCoJFnhFaHGyXlwWObjJyqnnocFoR1I2DSJtndm6Rc,2654
+timesfm/utils/xreg_lib.py,sha256=RgmzwRqhdE2zZcih9wf3Kkm2iUi8XWseDr2IZ3Kj9NI,20721
+tsagentkit_timesfm-1.0.0.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+tsagentkit_timesfm-1.0.0.dist-info/METADATA,sha256=_hzD2eY53jd9XY2CmlVlhqumxRLmp-Phu_C4nhNo1ks,4823
+tsagentkit_timesfm-1.0.0.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+tsagentkit_timesfm-1.0.0.dist-info/top_level.txt,sha256=5CtB5agsTzKTMKI3M87NsaJk7gh51fFKxmwRkdfeLdg,8
+tsagentkit_timesfm-1.0.0.dist-info/RECORD,,

tsagentkit_timesfm-1.0.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.10.2)
+Root-Is-Purelib: true
+Tag: py3-none-any
+