Repository: sktime/sktime

Issue:
[ENH] improvements for html repr of estimators
There are two major (but not breaking) issues with the current html representation of estimators, from `__repr__`:

1. Pipelines render left-to-right rather than top-to-bottom.

I am not sure how to fix this, the best would be to add a tag that somehow controls this for a given composite. See display below for the `ForecastingPipeline`.

```python
from sktime.forecasting.compose import ForecastingPipeline

ForecastingPipeline.create_test_instance()
```
yields 
![image](https://github.com/user-attachments/assets/51009021-6147-4a95-a9ac-6815273c125a)

2. The html display does not seem to correctly display dropdowns for multiple nested estimators across packages.

An example of this is the following estimator composite which has an sklearn estimator inside an skpro estimator, inside an sktime estimator:

```python
from sktime.forecasting.compose import DirectTabularRegressionForecaster
fs = DirectTabularRegressionForecaster.create_test_instances_and_names()
fs[0][2]
```

Interestingly, the sklearn estimator is displayed, but not the skpro estimator:
![image](https://github.com/user-attachments/assets/f9e94d03-e4a8-44d2-9271-983d1c307b98)

Discussion:
mateuszkasprowicz: Next step now would be to

1. Update https://github.com/sktime/sktime/blob/80c9eed218bac737dd4ca17a43d86a344236aaa4/docs/source/conf.py#L500-L511
2. Remove https://github.com/sktime/sktime/blob/80c9eed218bac737dd4ca17a43d86a344236aaa4/docs/source/estimator_overview.md?plain=1#L324-L337

And I believe this related to this https://github.com/sktime/sktime/issues/4905#issuecomment-2290981781

Any ideas how to deal with this at this time? 
fkiraly: Yes, I would consolidate this into *python* code, as we will need to call this from two locations - would that work?
mateuszkasprowicz: How about 
1. remove `go2URL` altogether
2. import `_get_reduced_path` from `sktime.utils._estimator_html_repr` to `conf.py`
3. create a simple tag navigating directly to correct url
```python
reduced_path = _get_reduced_path(modpath)
modname = (
    f"""<a href='api_reference/auto_generated/{reduced_path}.html'>"""
    f"""{obj_name}</a>"""
)
```
fkiraly: agreed
fkiraly: FYI @yarnabrina, @duydl, any thoughts?
mateuszkasprowicz: btw, what's `import_path` used for? do we need it?
https://github.com/sktime/sktime/blob/5ae6c637de8fef5dcd6f1fe3ac376f2d36ca109f/docs/source/conf.py#L534
fkiraly: don't know, maybe it is a vestigial remnant of an earlier attempt to fix the import path issue gracefully?
duydl: > FYI @yarnabrina, @duydl, any thoughts?

I want to note that the problem with the API path in the docs is not because of parsing `_` annotating, but rather with inconsistencies in when modules imported into `__init__`, which shortens the path by one element. While it is the convention to import objects from a "_private" file into `__init__`, sktime does not always follow this. This also causes the API summary table to be disordered. One example is `CNNClassifier`, the link is currently broken in this PR: https://sktime--6930.org.readthedocs.build/en/6930/api_reference/auto_generated/sktime.classification.deep_learning.cnn.CNNClassifier.html. 
I believe there would be a lot more. 
I actually tried the current approach in `_estimator_html_repr.py` before, but had to create `go2URL` to cover all cases.


> btw, what's `import_path` used for? do we need it?
> 
> https://github.com/sktime/sktime/blob/5ae6c637de8fef5dcd6f1fe3ac376f2d36ca109f/docs/source/conf.py#L534

This is simply for adding import paths to a data frame then json then finally list it in the estimator table. 

Patch:
diff --git a/docs/source/api_reference/forecasting.rst b/docs/source/api_reference/forecasting.rst
index 18efeceb8..02930a77f 100644
--- a/docs/source/api_reference/forecasting.rst
+++ b/docs/source/api_reference/forecasting.rst
@@ -459,6 +459,14 @@ Deep learning based forecasters
 
     HFTransformersForecaster
 
+.. currentmodule:: sktime.forecasting.timesfm_forecaster
+
+.. autosummary::
+    :toctree: auto_generated/
+    :template: class.rst
+
+    TimesFMForecaster
+
 .. currentmodule:: sktime.forecasting.ttm
 
 .. autosummary::
diff --git a/sktime/datatypes/_adapter/polars.py b/sktime/datatypes/_adapter/polars.py
index 257cb3652..e1fdd5f3a 100644
--- a/sktime/datatypes/_adapter/polars.py
+++ b/sktime/datatypes/_adapter/polars.py
@@ -29,7 +29,7 @@ def is_monotonically_increasing(obj):
         import polars as pl
 
         index_df = obj.with_columns(index_cols)
-        grouped = index_df.groupby(index_cols[:-1]).agg([pl.col(index_cols[-1])])
+        grouped = index_df.group_by(index_cols[:-1]).agg([pl.col(index_cols[-1])])
         last_index_col = grouped.select([index_cols[-1]])
         for val in last_index_col.iter_rows():
             # iter rows returns a list of tuples
@@ -252,14 +252,6 @@ def check_polars_frame(
             else:
                 metadata["n_instances"] = "NA"
 
-    if scitype in ["Panel", "Hierarchical"]:
-        if _req("n_instances", return_metadata):
-            if exp_type_str == "LazyFrame":
-                metadata["n_instances"] = "NA"
-            else:
-                instance_cols = index_cols[:-1]
-                metadata["n_instances"] = len(obj[instance_cols].unique())
-
     # check if there are any nans
     #   compute only if needed
     if _req("has_nans", return_metadata):
diff --git a/sktime/datatypes/_panel/_check.py b/sktime/datatypes/_panel/_check.py
index dfbca23f5..454b2bbcf 100644
--- a/sktime/datatypes/_panel/_check.py
+++ b/sktime/datatypes/_panel/_check.py
@@ -540,6 +540,81 @@ if _check_soft_dependencies("dask", severity="none"):
 
     check_dict[("dask_panel", "Panel")] = check_dask_panel
 
+if _check_soft_dependencies("polars", severity="none"):
+    import polars as pl
+
+    from sktime.datatypes._adapter.polars import check_polars_frame, get_mi_cols
+
+    def check_polars_panel(obj, return_metadata=False, var_name="obj", panel=True):
+        metadict = check_polars_frame(
+            obj=obj,
+            return_metadata=return_metadata,
+            var_name=var_name,
+            scitype="Panel",
+        )
+
+        if isinstance(metadict, tuple) and metadict[0]:
+            requires_series_grps = [
+                "n_instances",
+                "is_one_series",
+                "is_equal_length",
+                "is_equally_spaced",
+            ]
+
+            if _req(requires_series_grps, return_metadata):
+                if isinstance(obj, pl.LazyFrame):
+                    for key in requires_series_grps:
+                        metadict[1][key] = "NA"
+                else:
+                    mi_cols = get_mi_cols(obj)
+                    levels = mi_cols[:-1]
+                    if len(mi_cols) == 1:
+                        levels = levels[0]
+                    series_groups = obj.group_by(levels).agg(
+                        [
+                            pl.col(mi_cols[-1]).alias(
+                                "level_values"
+                            ),  # index values in each level
+                            pl.len().alias("size"),  # Calculate the size of each group
+                        ]
+                    )
+                    n_series = series_groups.shape[0]
+
+                    if _req("n_instances", return_metadata):
+                        metadict[2]["n_instances"] = n_series
+                    if _req("is_one_series", return_metadata):
+                        metadict[2]["is_one_series"] = n_series == 1
+                    if _req("is_equal_length", return_metadata):
+                        metadict[2]["is_equal_length"] = _list_all_equal(
+                            series_groups["size"].to_numpy()
+                        )
+                    if _req("is_equally_spaced", return_metadata):
+                        metadict[2]["is_equally_spaced"] = "NA"
+
+            requires_panel_grps = ["n_panels", "is_one_panel"]
+            if _req(requires_panel_grps, return_metadata):
+                if isinstance(obj, pl.LazyFrame):
+                    for key in requires_panel_grps:
+                        metadict[1][key] = "NA"
+                else:
+                    if panel:
+                        n_panels = 1
+                    else:
+                        mi_cols = get_mi_cols(obj)
+                        panel_groups = obj.group_by(mi_cols[:-2]).agg(
+                            [pl.col(mi_cols[-2])]
+                        )
+                        n_panels = panel_groups.shape[0]
+
+                    if _req("n_panels", return_metadata):
+                        metadict[2]["n_panels"] = n_panels
+                    if _req("is_one_panel", return_metadata):
+                        metadict[2]["is_one_panel"] = n_panels == 1
+
+        return metadict
+
+    check_dict[("polars_panel", "Panel")] = check_polars_panel
+
 if _check_soft_dependencies("gluonts", severity="none"):
     from sktime.datatypes._dtypekind import DtypeKind
 
diff --git a/sktime/datatypes/_panel/_convert.py b/sktime/datatypes/_panel/_convert.py
index c86fc29f8..81ff16e45 100644
--- a/sktime/datatypes/_panel/_convert.py
+++ b/sktime/datatypes/_panel/_convert.py
@@ -1127,6 +1127,30 @@ if _check_soft_dependencies("dask", severity="none"):
         "dask_panel", "pd-multiindex", convert_dict, mtype_universe=MTYPE_LIST_PANEL
     )
 
+if _check_soft_dependencies("polars", severity="none"):
+    from sktime.datatypes._adapter.polars import (
+        convert_pandas_to_polars,
+        convert_polars_to_pandas,
+    )
+
+    def convert_polars_to_pd_as_panel(obj, store=None):
+        return convert_polars_to_pandas(obj)
+
+    convert_dict[("polars_panel", "pd-multiindex", "Panel")] = (
+        convert_polars_to_pd_as_panel
+    )
+
+    def convert_pd_to_polars_as_panel(obj, store=None):
+        return convert_pandas_to_polars(obj)
+
+    convert_dict[("pd-multiindex", "polars_panel", "Panel")] = (
+        convert_pd_to_polars_as_panel
+    )
+
+    _extend_conversions(
+        "polars_panel", "pd-multiindex", convert_dict, mtype_universe=MTYPE_LIST_PANEL
+    )
+
 if _check_soft_dependencies("gluonts", severity="none"):
     from sktime.datatypes._adapter.gluonts import (
         convert_listDataset_to_pandas,
diff --git a/sktime/datatypes/_panel/_examples.py b/sktime/datatypes/_panel/_examples.py
index a4e8aa998..242bf34c7 100644
--- a/sktime/datatypes/_panel/_examples.py
+++ b/sktime/datatypes/_panel/_examples.py
@@ -97,6 +97,14 @@ if _check_soft_dependencies("dask", severity="none"):
     example_dict[("dask_panel", "Panel", 0)] = df_dask
     example_dict_lossy[("dask_panel", "Panel", 0)] = False
 
+if _check_soft_dependencies("polars", severity="none"):
+    from sktime.datatypes._adapter.polars import convert_pandas_to_polars
+
+    pl_frame = convert_pandas_to_polars(example_dict[("pd-multiindex", "Panel", 0)])
+
+    example_dict[("polars_panel", "Panel", 0)] = pl_frame
+    example_dict_lossy[("polars_panel", "Panel", 0)] = False
+
 if _check_soft_dependencies("gluonts", severity="none"):
     from sktime.datatypes._adapter.gluonts import (
         convert_pandas_multiindex_to_pandasDataset,
@@ -199,6 +207,13 @@ if _check_soft_dependencies("dask", severity="none"):
     example_dict[("dask_panel", "Panel", 1)] = df_dask
     example_dict_lossy[("dask_panel", "Panel", 1)] = False
 
+if _check_soft_dependencies("polars", severity="none"):
+    from sktime.datatypes._adapter.polars import convert_pandas_to_polars
+
+    pl_frame = convert_pandas_to_polars(example_dict[("pd-multiindex", "Panel", 1)])
+
+    example_dict[("polars_panel", "Panel", 1)] = pl_frame
+    example_dict_lossy[("polars_panel", "Panel", 1)] = False
 if _check_soft_dependencies("gluonts", severity="none"):
     from sktime.datatypes._adapter.gluonts import convert_pandas_to_listDataset
 
@@ -290,6 +305,14 @@ if _check_soft_dependencies("dask", severity="none"):
     example_dict[("dask_panel", "Panel", 2)] = df_dask
     example_dict_lossy[("dask_panel", "Panel", 2)] = False
 
+if _check_soft_dependencies("polars", severity="none"):
+    from sktime.datatypes._adapter.polars import convert_pandas_to_polars
+
+    pl_frame = convert_pandas_to_polars(example_dict[("pd-multiindex", "Panel", 2)])
+
+    example_dict[("polars_panel", "Panel", 2)] = pl_frame
+    example_dict_lossy[("polars_panel", "Panel", 2)] = False
+
 if _check_soft_dependencies("gluonts", severity="none"):
     from sktime.datatypes._adapter.gluonts import convert_pandas_to_listDataset
 
diff --git a/sktime/datatypes/_panel/_registry.py b/sktime/datatypes/_panel/_registry.py
index 077d10a1d..344326eef 100644
--- a/sktime/datatypes/_panel/_registry.py
+++ b/sktime/datatypes/_panel/_registry.py
@@ -45,6 +45,11 @@ MTYPE_REGISTER_PANEL = [
         "Panel",
         "gluonTS representation of a pandas DataFrame",
     ),
+    (
+        "polars_panel",
+        "Panel",
+        "Polars frame representing a Panel, with one instance and one time index",
+    ),
 ]
 
 MTYPE_SOFT_DEPS_PANEL = {
@@ -52,6 +57,7 @@ MTYPE_SOFT_DEPS_PANEL = {
     "dask_panel": "dask",
     "gluonts_ListDataset_panel": "gluonts",
     "gluonts_PandasDataset_panel": "gluonts",
+    "polars_panel": "polars",
 }
 
 MTYPE_LIST_PANEL = [x[0] for x in MTYPE_REGISTER_PANEL]
diff --git a/sktime/forecasting/compose/_pipeline.py b/sktime/forecasting/compose/_pipeline.py
index e89cd9aa0..bd80f57c0 100644
--- a/sktime/forecasting/compose/_pipeline.py
+++ b/sktime/forecasting/compose/_pipeline.py
@@ -13,6 +13,7 @@ from sktime.forecasting.base._base import BaseForecaster
 from sktime.forecasting.base._delegate import _DelegatedForecaster
 from sktime.forecasting.base._fh import ForecastingHorizon
 from sktime.registry import scitype
+from sktime.utils._estimator_html_repr import _VisualBlock
 from sktime.utils.validation.series import check_series
 from sktime.utils.warnings import warn
 
@@ -278,6 +279,25 @@ class _Pipeline(_HeterogenousMetaEstimator, BaseForecaster):
 
         return [params1, params2, params3]
 
+    def _sk_visual_block_(self):
+        _, estimators = zip(*self.steps)
+
+        def _get_name(name, est):
+            if est is None or est == "passthrough":
+                return f"{name}: passthrough"
+            # Is an estimator
+            return f"{name}: {est.__class__.__name__}"
+
+        names = [_get_name(name, est) for name, est in self.steps]
+        name_details = [str(est) for est in estimators]
+        return _VisualBlock(
+            "serial",
+            estimators,
+            names=names,
+            name_details=name_details,
+            dash_wrapped=False,
+        )
+
 
 # we ensure that internally we convert to pd.DataFrame for now
 SUPPORTED_MTYPES = ["pd.DataFrame", "pd-multiindex", "pd_multiindex_hier"]
diff --git a/sktime/forecasting/timesfm_forecaster.py b/sktime/forecasting/timesfm_forecaster.py
new file mode 100644
index 000000000..35fb3ea69
--- /dev/null
+++ b/sktime/forecasting/timesfm_forecaster.py
@@ -0,0 +1,345 @@
+# copyright: sktime developers, BSD-3-Clause License (see LICENSE file)
+"""Implementation of TimesFM (Time Series Foundation Model)."""
+
+__author__ = ["rajatsen91", "geetu040"]
+# rajatsen91 for google-research/timesfm
+
+import os
+
+import numpy as np
+import pandas as pd
+
+from sktime.forecasting.base import ForecastingHorizon, _BaseGlobalForecaster
+
+
+class TimesFMForecaster(_BaseGlobalForecaster):
+    """
+    Implementation of TimesFM (Time Series Foundation Model) for Zero-Shot Forecasting.
+
+    TimesFM (Time Series Foundation Model) is a pretrained time-series foundation model
+    developed by Google Research for time-series forecasting. This method has been
+    proposed in [2]_ and official code is given at [1]_.
+
+    Parameters
+    ----------
+    context_len : int
+        The length of the input context sequence.
+        It should be a multiple of ``input_patch_len`` (32).
+        The maximum context length currently supported is 512, but this can be
+        increased in future releases.
+        The input time series can have any context length,
+        and padding or truncation will
+        be handled by the model's inference code if necessary.
+    horizon_len : int
+        The length of the forecast horizon. This can be set to any value, although it
+        is generally recommended to keep it less than or equal to ``context_len`` for
+        optimal performance. The model will still function
+        if ``horizon_len`` exceeds ``context_len``.
+    freq : int, optional (default=0)
+        The frequency category of the input time series.
+        - 0: High frequency, long horizon time series (e.g., daily data and above).
+        - 1: Medium frequency time series (e.g., weekly, monthly data).
+        - 2: Low frequency, short horizon time series (e.g., quarterly, yearly data).
+        You can treat this parameter as a free parameter depending
+        on your specific use case,
+        although it is recommended to follow these guidelines for optimal results.
+    repo_id : str, optional (default="google/timesfm-1.0-200m")
+        The identifier for the model repository.
+        The default model is the 200M parameter version.
+    input_patch_len : int, optional (default=32)
+        The fixed length of input patches that the model processes.
+        This parameter is fixed to 1280 for the 200M model and should not be changed.
+    output_patch_len : int, optional (default=128)
+        The fixed length of output patches that the model generates.
+        This parameter is fixed to 1280 for the 200M model and should not be changed.
+    num_layers : int, optional (default=20)
+        The number of layers in the model architecture.
+        This parameter is fixed to 1280 for the 200M model and should not be changed.
+    model_dims : int, optional (default=1280)
+        The dimensionality of the model.
+        This parameter is fixed to 1280 for the 200M model and should not be changed.
+    per_core_batch_size : int, optional (default=32)
+        The batch size to be used per core during model inference.
+    backend : str, optional (default="cpu")
+        The computational backend to be used,
+        which can be one of "cpu", "gpu", or "tpu".
+        This setting is case-sensitive.
+    verbose : bool, optional (default=False)
+        Whether to print detailed logs during execution.
+    broadcasting : bool, default=False
+        if True, multiindex data input will be broadcasted to single series.
+        For each single series, one copy of this forecaster will try to
+        fit and predict on it. The broadcasting is happening inside automatically,
+        from the outerside api perspective, the input and output are the same,
+        only one multiindex output from ``predict``.
+    use_source_package : bool, default=False
+        If True, the model will be loaded directly from the source package ``timesfm``.
+        This is useful if you want to bypass the local version of the package
+        or when working in an environment where the latest updates
+        from the source package are needed.
+        If False, the model will be loaded from the local version of package maintained
+        in sktime.
+        To install the source package, follow the instructions here [1]_.
+
+    References
+    ----------
+    .. [1] https://github.com/google-research/timesfm
+    .. [2] Das, A., Kong, W., Sen, R., & Zhou, Y. (2024).
+    A decoder-only foundation model for time-series forecasting. CoRR.
+
+    Examples
+    --------
+    >>> from sktime.forecasting.timesfm_forecaster import TimesFMForecaster
+    >>> from sktime.datasets import load_airline
+    >>> y = load_airline()
+    >>> forecaster = TimesFMForecaster(
+    ...     context_len=32,
+    ...     horizon_len=8,
+    ... ) # doctest: +SKIP
+    >>> forecaster.fit(y, fh=[1, 2, 3]) # doctest: +SKIP
+    >>> y_pred = forecaster.predict() # doctest: +SKIP
+
+    >>> from sktime.forecasting.timesfm_forecaster import TimesFMForecaster
+    >>> from sktime.datasets import load_tecator
+    >>>
+    >>> # load multi-index dataset
+    >>> y = load_tecator(
+    ...     return_type="pd-multiindex",
+    ...     return_X_y=False
+    ... )
+    >>> y.drop(['class_val'], axis=1, inplace=True)
+    >>>
+    >>> # global forecasting on multi-index dataset
+    >>> forecaster = TimesFMForecaster(
+    ...     context_len=32,
+    ...     horizon_len=8,
+    ... ) # doctest: +SKIP
+    >>>
+    >>> # train and predict
+    >>> forecaster.fit(y, fh=[1, 2, 3]) # doctest: +SKIP
+    >>> y_pred = forecaster.predict() # doctest: +SKIP
+    """
+
+    _tags = {
+        "y_inner_mtype": [
+            "pd.Series",
+            "pd-multiindex",
+            "pd_multiindex_hier",
+        ],
+        "scitype:y": "univariate",
+        "ignores-exogeneous-X": True,
+        "requires-fh-in-fit": False,
+        "X-y-must-have-same-index": True,
+        "enforce_index_type": None,
+        "handles-missing-data": False,
+        "capability:insample": False,
+        "capability:pred_int": False,
+        "capability:pred_int:insample": False,
+        "authors": ["rajatsen91", "geetu040"],
+        # rajatsen91 for google-research/timesfm
+        "maintainers": ["geetu040"],
+        "python_version": ">=3.10,<3.11",
+        "python_dependencies": [
+            "tensorflow",
+            "einshape",
+            "jax",
+            "praxis",
+            "huggingface-hub",
+            "paxml",
+            "utilsforecast",
+        ],
+        "env_marker": "sys_platform=='linux'",
+        "capability:global_forecasting": True,
+    }
+
+    def __init__(
+        self,
+        context_len,
+        horizon_len,
+        freq=0,
+        repo_id="google/timesfm-1.0-200m",
+        input_patch_len=32,
+        output_patch_len=128,
+        num_layers=20,
+        model_dims=1280,
+        per_core_batch_size=32,
+        backend="cpu",
+        verbose=False,
+        broadcasting=False,
+        use_source_package=False,
+    ):
+        self.context_len = context_len
+        self.horizon_len = horizon_len
+        self._horizon_len = None
+        self.freq = freq
+        self.repo_id = repo_id
+        self.input_patch_len = input_patch_len
+        self.output_patch_len = output_patch_len
+        self.num_layers = num_layers
+        self.model_dims = model_dims
+        self.per_core_batch_size = per_core_batch_size
+        self.backend = backend
+        self.verbose = verbose
+        self.broadcasting = broadcasting
+        self.use_source_package = use_source_package
+
+        if self.broadcasting:
+            self.set_tags(
+                **{
+                    "y_inner_mtype": "pd.Series",
+                    "X_inner_mtype": "pd.DataFrame",
+                    "capability:global_forecasting": False,
+                }
+            )
+
+        # to avoid RuntimeError when backed=="cpu"
+        os.environ["JAX_PLATFORM_NAME"] = backend
+        os.environ["JAX_PLATFORMS"] = backend
+
+        super().__init__()
+
+    def _fit(self, y, X, fh):
+        # import after backend env has been set
+        if self.use_source_package:
+            from timesfm import TimesFm
+        else:
+            from sktime.libs.timesfm import TimesFm
+
+        if fh is not None:
+            fh = fh.to_relative(self.cutoff)
+            self._horizon_len = max(self.horizon_len, *fh._values.values)
+        else:
+            self._horizon_len = self.horizon_len
+
+        self.tfm = TimesFm(
+            context_len=self.context_len,
+            horizon_len=self._horizon_len,
+            input_patch_len=self.input_patch_len,
+            output_patch_len=self.output_patch_len,
+            num_layers=self.num_layers,
+            model_dims=self.model_dims,
+            per_core_batch_size=self.per_core_batch_size,
+            backend=self.backend,
+            verbose=self.verbose,
+        )
+        self.tfm.load_from_checkpoint(repo_id=self.repo_id)
+
+    def _predict(self, fh, X, y=None):
+        if fh is None:
+            fh = self.fh
+        fh = fh.to_relative(self.cutoff)
+
+        if max(fh._values.values) > self.horizon_len:
+            raise ValueError(
+                f"Error in {self.__class__.__name__}, the forecast horizon exceeds the"
+                f" specified horizon_len of {self.horizon_len}. Change the horizon_len"
+                " when initializing the model or try another forecasting horizon."
+            )
+
+        _y = y if self._global_forecasting else self._y
+
+        # multi-index conversion goes here
+        if isinstance(_y.index, pd.MultiIndex):
+            hist = _frame2numpy(_y).squeeze(2)
+        else:
+            hist = np.expand_dims(_y.values, axis=0)
+
+        # hist.shape: (batch_size, n_timestamps)
+
+        pred, _ = self.tfm.forecast(hist)
+
+        # converting pred datatype
+
+        batch_size, n_timestamps = pred.shape
+
+        if isinstance(_y.index, pd.MultiIndex):
+            ins = np.array(list(np.unique(_y.index.droplevel(-1)).repeat(n_timestamps)))
+            ins = [ins[..., i] for i in range(ins.shape[-1])] if ins.ndim > 1 else [ins]
+
+            idx = (
+                ForecastingHorizon(range(1, n_timestamps + 1), freq=self.fh.freq)
+                .to_absolute(self._cutoff)
+                ._values.tolist()
+                * pred.shape[0]
+            )
+            index = pd.MultiIndex.from_arrays(
+                ins + [idx],
+                names=_y.index.names,
+            )
+            pred = pd.DataFrame(
+                # batch_size * num_timestams
+                pred.ravel(),
+                index=index,
+                columns=_y.columns,
+            )
+        else:
+            index = (
+                ForecastingHorizon(range(1, n_timestamps + 1))
+                .to_absolute(self._cutoff)
+                ._values
+            )
+            pred = pd.Series(
+                # batch_size * num_timestams
+                pred.ravel(),
+                index=index,
+                name=_y.name,
+            )
+
+        absolute_horizons = fh.to_absolute_index(self.cutoff)
+        dateindex = pred.index.get_level_values(-1).map(
+            lambda x: x in absolute_horizons
+        )
+        pred = pred.loc[dateindex]
+
+        return pred
+
+    @classmethod
+    def get_test_params(cls, parameter_set="default"):
+        """Return testing parameter settings for the estimator.
+
+        Parameters
+        ----------
+        parameter_set : str, default="default"
+            Name of the set of test parameters to return, for use in tests. If no
+            special parameters are defined for a value, will return `"default"` set.
+            There are currently no reserved values for forecasters.
+
+        Returns
+        -------
+        params : dict or list of dict, default = {}
+            Parameters to create testing instances of the class
+            Each dict are parameters to construct an "interesting" test instance, i.e.,
+            `MyClass(**params)` or `MyClass(**params[i])` creates a valid test instance.
+            `create_test_instance` uses the first (or only) dictionary in `params`
+        """
+        test_params = [
+            {
+                "context_len": 64,
+                "horizon_len": 32,
+                "freq": 0,
+                "verbose": False,
+            },
+        ]
+        params_no_broadcasting = [
+            dict(p, **{"broadcasting": True}) for p in test_params
+        ]
+        test_params.extend(params_no_broadcasting)
+        return test_params
+
+
+def _same_index(data):
+    data = data.groupby(level=list(range(len(data.index.levels) - 1))).apply(
+        lambda x: x.index.get_level_values(-1)
+    )
+    assert data.map(
+        lambda x: x.equals(data.iloc[0])
+    ).all(), "All series must has the same index"
+    return data.iloc[0], len(data.iloc[0])
+
+
+def _frame2numpy(data):
+    idx, length = _same_index(data)
+    arr = np.array(data.values, dtype=np.float32).reshape(
+        (-1, length, len(data.columns))
+    )
+    return arr
diff --git a/sktime/libs/timesfm/__init__.py b/sktime/libs/timesfm/__init__.py
new file mode 100644
index 000000000..48dbf5f87
--- /dev/null
+++ b/sktime/libs/timesfm/__init__.py
@@ -0,0 +1,220 @@
+"""Python implementation of TimesFM (Time Series Foundation Model).
+
+Unofficial fork of the ``google-research/timesfm`` package, maintained in ``sktime``.
+This implementation is intended as a temporary solution, until the authors
+release a package that resolves dependency problems as raised
+here https://github.com/google-research/timesfm/issues/1.
+
+sktime migration: 2024, Aug
+Version 1.0.1 release: 2024, Aug (rajatsen91)
+
+Original author: Rajat Sen
+
+2024 releases subject to following license:
+
+All code contained is released under the license below.
+
+Apache License
+Version 2.0, January 2004
+http://www.apache.org/licenses/
+
+TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+1. Definitions.
+
+        "License" shall mean the terms and conditions for use, reproduction,
+        and distribution as defined by Sections 1 through 9 of this document.
+
+        "Licensor" shall mean the copyright owner or entity authorized by
+        the copyright owner that is granting the License.
+
+        "Legal Entity" shall mean the union of the acting entity and all
+        other entities that control, are controlled by, or are under common
+        control with that entity. For the purposes of this definition,
+        "control" means (i) the power, direct or indirect, to cause the
+        direction or management of such entity, whether by contract or
+        otherwise, or (ii) ownership of fifty percent (50%) or more of the
+        outstanding shares, or (iii) beneficial ownership of such entity.
+
+        "You" (or "Your") shall mean an individual or Legal Entity
+        exercising permissions granted by this License.
+
+        "Source" form shall mean the preferred form for making modifications,
+        including but not limited to software source code, documentation
+        source, and configuration files.
+
+        "Object" form shall mean any form resulting from mechanical
+        transformation or translation of a Source form, including but
+        not limited to compiled object code, generated documentation,
+        and conversions to other media types.
+
+        "Work" shall mean the work of authorship, whether in Source or
+        Object form, made available under the License, as indicated by a
+        copyright notice that is included in or attached to the work
+        (an example is provided in the Appendix below).
+
+        "Derivative Works" shall mean any work, whether in Source or Object
+        form, that is based on (or derived from) the Work and for which the
+        editorial revisions, annotations, elaborations, or other modifications
+        represent, as a whole, an original work of authorship. For the purposes
+        of this License, Derivative Works shall not include works that remain
+        separable from, or merely link (or bind by name) to the interfaces of,
+        the Work and Derivative Works thereof.
+
+        "Contribution" shall mean any work of authorship, including
+        the original version of the Work and any modifications or additions
+        to that Work or Derivative Works thereof, that is intentionally
+        submitted to Licensor for inclusion in the Work by the copyright owner
+        or by an individual or Legal Entity authorized to submit on behalf of
+        the copyright owner. For the purposes of this definition, "submitted"
+        means any form of electronic, verbal, or written communication sent
+        to the Licensor or its representatives, including but not limited to
+        communication on electronic mailing lists, source code control systems,
+        and issue tracking systems that are managed by, or on behalf of, the
+        Licensor for the purpose of discussing and improving the Work, but
+        excluding communication that is conspicuously marked or otherwise
+        designated in writing by the copyright owner as "Not a Contribution."
+
+        "Contributor" shall mean Licensor and any individual or Legal Entity
+        on behalf of whom a Contribution has been received by Licensor and
+        subsequently incorporated within the Work.
+
+2. Grant of Copyright License. Subject to the terms and conditions of
+        this License, each Contributor hereby grants to You a perpetual,
+        worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+        copyright license to reproduce, prepare Derivative Works of,
+        publicly display, publicly perform, sublicense, and distribute the
+        Work and such Derivative Works in Source or Object form.
+
+3. Grant of Patent License. Subject to the terms and conditions of
+        this License, each Contributor hereby grants to You a perpetual,
+        worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+        (except as stated in this section) patent license to make, have made,
+        use, offer to sell, sell, import, and otherwise transfer the Work,
+        where such license applies only to those patent claims licensable
+        by such Contributor that are necessarily infringed by their
+        Contribution(s) alone or by combination of their Contribution(s)
+        with the Work to which such Contribution(s) was submitted. If You
+        institute patent litigation against any entity (including a
+        cross-claim or counterclaim in a lawsuit) alleging that the Work
+        or a Contribution incorporated within the Work constitutes direct
+        or contributory patent infringement, then any patent licenses
+        granted to You under this License for that Work shall terminate
+        as of the date such litigation is filed.
+
+4. Redistribution. You may reproduce and distribute copies of the
+        Work or Derivative Works thereof in any medium, with or without
+        modifications, and in Source or Object form, provided that You
+        meet the following conditions:
+
+        (a) You must give any other recipients of the Work or
+                Derivative Works a copy of this License; and
+
+        (b) You must cause any modified files to carry prominent notices
+                stating that You changed the files; and
+
+        (c) You must retain, in the Source form of any Derivative Works
+                that You distribute, all copyright, patent, trademark, and
+                attribution notices from the Source form of the Work,
+                excluding those notices that do not pertain to any part of
+                the Derivative Works; and
+
+        (d) If the Work includes a "NOTICE" text file as part of its
+                distribution, then any Derivative Works that You distribute must
+                include a readable copy of the attribution notices contained
+                within such NOTICE file, excluding those notices that do not
+                pertain to any part of the Derivative Works, in at least one
+                of the following places: within a NOTICE text file distributed
+                as part of the Derivative Works; within the Source form or
+                documentation, if provided along with the Derivative Works; or,
+                within a display generated by the Derivative Works, if and
+                wherever such third-party notices normally appear. The contents
+                of the NOTICE file are for informational purposes only and
+                do not modify the License. You may add Your own attribution
+                notices within Derivative Works that You distribute, alongside
+                or as an addendum to the NOTICE text from the Work, provided
+                that such additional attribution notices cannot be construed
+                as modifying the License.
+
+        You may add Your own copyright statement to Your modifications and
+        may provide additional or different license terms and conditions
+        for use, reproduction, or distribution of Your modifications, or
+        for any such Derivative Works as a whole, provided Your use,
+        reproduction, and distribution of the Work otherwise complies with
+        the conditions stated in this License.
+
+5. Submission of Contributions. Unless You explicitly state otherwise,
+        any Contribution intentionally submitted for inclusion in the Work
+        by You to the Licensor shall be under the terms and conditions of
+        this License, without any additional terms or conditions.
+        Notwithstanding the above, nothing herein shall supersede or modify
+        the terms of any separate license agreement you may have executed
+        with Licensor regarding such Contributions.
+
+6. Trademarks. This License does not grant permission to use the trade
+        names, trademarks, service marks, or product names of the Licensor,
+        except as required for reasonable and customary use in describing the
+        origin of the Work and reproducing the content of the NOTICE file.
+
+7. Disclaimer of Warranty. Unless required by applicable law or
+        agreed to in writing, Licensor provides the Work (and each
+        Contributor provides its Contributions) on an "AS IS" BASIS,
+        WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+        implied, including, without limitation, any warranties or conditions
+        of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+        PARTICULAR PURPOSE. You are solely responsible for determining the
+        appropriateness of using or redistributing the Work and assume any
+        risks associated with Your exercise of permissions under this License.
+
+8. Limitation of Liability. In no event and under no legal theory,
+        whether in tort (including negligence), contract, or otherwise,
+        unless required by applicable law (such as deliberate and grossly
+        negligent acts) or agreed to in writing, shall any Contributor be
+        liable to You for damages, including any direct, indirect, special,
+        incidental, or consequential damages of any character arising as a
+        result of this License or out of the use or inability to use the
+        Work (including but not limited to damages for loss of goodwill,
+        work stoppage, computer failure or malfunction, or any and all
+        other commercial damages or losses), even if such Contributor
+        has been advised of the possibility of such damages.
+
+9. Accepting Warranty or Additional Liability. While redistributing
+        the Work or Derivative Works thereof, You may choose to offer,
+        and charge a fee for, acceptance of support, warranty, indemnity,
+        or other liability obligations and/or rights consistent with this
+        License. However, in accepting such obligations, You may act only
+        on Your own behalf and on Your sole responsibility, not on behalf
+        of any other Contributor, and only if You agree to indemnify,
+        defend, and hold each Contributor harmless for any liability
+        incurred by, or claims asserted against, such Contributor by reason
+        of your accepting any such warranty or additional liability.
+
+END OF TERMS AND CONDITIONS
+
+APPENDIX: How to apply the Apache License to your work.
+
+        To apply the Apache License to your work, attach the following
+        boilerplate notice, with the fields enclosed by brackets "[]"
+        replaced with your own identifying information. (Don't include
+        the brackets!)  The text should be enclosed in the appropriate
+        comment syntax for the file format. We also recommend that a
+        file or class name and description of purpose be included on the
+        same "printed page" as the copyright notice for easier
+        identification within third-party archives.
+
+Copyright [yyyy] [name of copyright owner]
+
+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.
+"""
+
+from sktime.libs.timesfm.timesfm import TimesFm
diff --git a/sktime/libs/timesfm/data_loader.py b/sktime/libs/timesfm/data_loader.py
new file mode 100644
index 000000000..4ca37422f
--- /dev/null
+++ b/sktime/libs/timesfm/data_loader.py
@@ -0,0 +1,240 @@
+# Copyright 2024 The Google Research Authors.
+#
+# 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.
+
+"""data_loader."""
+
+import numpy as np
+import pandas as pd
+from sklearn.preprocessing import StandardScaler
+
+from sktime.libs.timesfm import time_features
+from sktime.utils.dependencies import _check_soft_dependencies
+
+if _check_soft_dependencies("tensorflow", severity="none"):
+    import tensorflow as tf
+
+# avoid another dependency
+# from absl import logging
+
+
+class TimeSeriesdata:
+    """TimeSeriesdata."""
+
+    def __init__(
+        self,
+        data_path,
+        datetime_col,
+        num_cov_cols,
+        cat_cov_cols,
+        ts_cols,
+        train_range,
+        val_range,
+        test_range,
+        hist_len,
+        pred_len,
+        batch_size,
+        freq="H",
+        normalize=True,
+        epoch_len=None,
+        holiday=False,
+        permute=True,
+    ):
+        """__init__."""
+        self.data_df = pd.read_csv(open(data_path))
+        if not num_cov_cols:
+            self.data_df["ncol"] = np.zeros(self.data_df.shape[0])
+            num_cov_cols = ["ncol"]
+        if not cat_cov_cols:
+            self.data_df["ccol"] = np.zeros(self.data_df.shape[0])
+            cat_cov_cols = ["ccol"]
+        self.data_df.fillna(0, inplace=True)
+        self.data_df.set_index(
+            pd.DatetimeIndex(self.data_df[datetime_col]), inplace=True
+        )
+        self.num_cov_cols = num_cov_cols
+        self.cat_cov_cols = cat_cov_cols
+        self.ts_cols = ts_cols
+        self.train_range = train_range
+        self.val_range = val_range
+        self.test_range = test_range
+        data_df_idx = self.data_df.index
+        date_index = data_df_idx.union(
+            pd.date_range(
+                data_df_idx[-1] + pd.Timedelta(1, freq=freq),
+                periods=pred_len + 1,
+                freq=freq,
+            )
+        )
+        self.time_df = time_features.TimeCovariates(
+            date_index, holiday=holiday
+        ).get_covariates()
+        self.hist_len = hist_len
+        self.pred_len = pred_len
+        self.batch_size = batch_size
+        self.freq = freq
+        self.normalize = normalize
+        self.data_mat = self.data_df[self.ts_cols].to_numpy().transpose()
+        self.data_mat = self.data_mat[:, 0 : self.test_range[1]]
+        self.time_mat = self.time_df.to_numpy().transpose()
+        self.num_feat_mat = self.data_df[num_cov_cols].to_numpy().transpose()
+        self.cat_feat_mat, self.cat_sizes = self._get_cat_cols(cat_cov_cols)
+        self.normalize = normalize
+        if normalize:
+            self._normalize_data()
+        # logging.info(
+        #     "Data Shapes: %s, %s, %s, %s",
+        #     self.data_mat.shape,
+        #     self.time_mat.shape,
+        #     self.num_feat_mat.shape,
+        #     self.cat_feat_mat.shape,
+        # )
+        self.epoch_len = epoch_len
+        self.permute = permute
+
+    def _get_cat_cols(self, cat_cov_cols):
+        cat_vars = []
+        cat_sizes = []
+        for col in cat_cov_cols:
+            dct = {x: i for i, x in enumerate(self.data_df[col].unique())}
+            cat_sizes.append(len(dct))
+            mapped = self.data_df[col].map(lambda x: dct[x]).to_numpy().transpose()
+            cat_vars.append(mapped)
+        return np.vstack(cat_vars), cat_sizes
+
+    def _normalize_data(self):
+        self.scaler = StandardScaler()
+        train_mat = self.data_mat[:, 0 : self.train_range[1]]
+        self.scaler = self.scaler.fit(train_mat.transpose())
+        self.data_mat = self.scaler.transform(self.data_mat.transpose()).transpose()
+
+    def train_gen(self):
+        """train_gen."""
+        num_ts = len(self.ts_cols)
+        perm = np.arange(
+            self.train_range[0] + self.hist_len,
+            self.train_range[1] - self.pred_len,
+        )
+        perm = np.random.permutation(perm)
+        hist_len = self.hist_len
+        # logging.info("Hist len: %s", hist_len)
+        if not self.epoch_len:
+            epoch_len = len(perm)
+        else:
+            epoch_len = self.epoch_len
+        for idx in perm[0:epoch_len]:
+            for _ in range(num_ts // self.batch_size + 1):
+                if self.permute:
+                    tsidx = np.random.choice(
+                        num_ts, size=self.batch_size, replace=False
+                    )
+                else:
+                    tsidx = np.arange(num_ts)
+                dtimes = np.arange(idx - hist_len, idx + self.pred_len)
+                (
+                    bts_train,
+                    bts_pred,
+                    bfeats_train,
+                    bfeats_pred,
+                    bcf_train,
+                    bcf_pred,
+                ) = self._get_features_and_ts(dtimes, tsidx, hist_len)
+                all_data = [
+                    bts_train,
+                    bfeats_train,
+                    bcf_train,
+                    bts_pred,
+                    bfeats_pred,
+                    bcf_pred,
+                    tsidx,
+                ]
+                yield tuple(all_data)
+
+    def test_val_gen(self, mode="val", shift=1):
+        """test_val_gen."""
+        if mode == "val":
+            start = self.val_range[0]
+            end = self.val_range[1] - self.pred_len + 1
+        elif mode == "test":
+            start = self.test_range[0]
+            end = self.test_range[1] - self.pred_len + 1
+        else:
+            raise NotImplementedError("Eval mode not implemented")
+        num_ts = len(self.ts_cols)
+        hist_len = self.hist_len
+        # logging.info("Hist len: %s", hist_len)
+        perm = np.arange(start, end)
+        if self.epoch_len:
+            epoch_len = self.epoch_len
+        else:
+            epoch_len = len(perm)
+        for i in range(0, epoch_len, shift):
+            idx = perm[i]
+            for batch_idx in range(0, num_ts, self.batch_size):
+                tsidx = np.arange(batch_idx, min(batch_idx + self.batch_size, num_ts))
+                dtimes = np.arange(idx - hist_len, idx + self.pred_len)
+                (
+                    bts_train,
+                    bts_pred,
+                    bfeats_train,
+                    bfeats_pred,
+                    bcf_train,
+                    bcf_pred,
+                ) = self._get_features_and_ts(dtimes, tsidx, hist_len)
+                all_data = [
+                    bts_train,
+                    bfeats_train,
+                    bcf_train,
+                    bts_pred,
+                    bfeats_pred,
+                    bcf_pred,
+                    tsidx,
+                ]
+                yield tuple(all_data)
+
+    def _get_features_and_ts(self, dtimes, tsidx, hist_len=None):
+        if hist_len is None:
+            hist_len = self.hist_len
+        data_times = dtimes[dtimes < self.data_mat.shape[1]]
+        bdata = self.data_mat[:, data_times]
+        bts = bdata[tsidx, :]
+        bnf = self.num_feat_mat[:, data_times]
+        bcf = self.cat_feat_mat[:, data_times]
+        btf = self.time_mat[:, dtimes]
+        if bnf.shape[1] < btf.shape[1]:
+            rem_len = btf.shape[1] - bnf.shape[1]
+            rem_rep = np.repeat(bnf[:, [-1]], repeats=rem_len)
+            rem_rep_cat = np.repeat(bcf[:, [-1]], repeats=rem_len)
+            bnf = np.hstack([bnf, rem_rep.reshape(bnf.shape[0], -1)])
+            bcf = np.hstack([bcf, rem_rep_cat.reshape(bcf.shape[0], -1)])
+        bfeats = np.vstack([btf, bnf])
+        bts_train = bts[:, 0:hist_len]
+        bts_pred = bts[:, hist_len:]
+        bfeats_train = bfeats[:, 0:hist_len]
+        bfeats_pred = bfeats[:, hist_len:]
+        bcf_train = bcf[:, 0:hist_len]
+        bcf_pred = bcf[:, hist_len:]
+        return bts_train, bts_pred, bfeats_train, bfeats_pred, bcf_train, bcf_pred
+
+    def tf_dataset(self, mode="train", shift=1):
+        """Tensorflow Dataset."""
+        if mode == "train":
+            gen_fn = self.train_gen
+        else:
+            gen_fn = lambda: self.test_val_gen(mode, shift)
+        output_types = tuple(
+            [tf.float32] * 2 + [tf.int32] + [tf.float32] * 2 + [tf.int32] * 2
+        )
+        dataset = tf.data.Dataset.from_generator(gen_fn, output_types)
+        dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
+        return dataset
diff --git a/sktime/libs/timesfm/patched_decoder.py b/sktime/libs/timesfm/patched_decoder.py
new file mode 100644
index 000000000..06f89105f
--- /dev/null
+++ b/sktime/libs/timesfm/patched_decoder.py
@@ -0,0 +1,477 @@
+# Copyright 2024 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.
+
+"""patched_decoder."""
+
+import dataclasses
+
+from sktime.utils.dependencies import _check_soft_dependencies
+
+if _check_soft_dependencies("einshape", severity="none"):
+    import einshape as es
+
+if _check_soft_dependencies("jax", severity="none"):
+    import jax.numpy as jnp
+    from jax import lax
+
+if _check_soft_dependencies("praxis", severity="none"):
+    from praxis import base_layer, layers, pax_fiddle, py_utils
+    from praxis.base_layer import BaseLayer
+    from praxis.base_model import BaseModel
+    from praxis.layers import (
+        activations,
+        embedding_softmax,
+        linears,
+        normalizations,
+        stochastics,
+        transformers,
+    )
+
+    # PAX shortcuts
+    NestedMap = py_utils.NestedMap
+
+    LayerTpl = pax_fiddle.Config[BaseLayer]
+    template_field = base_layer.template_field
+
+else:
+
+    class BaseLayer:
+        """Dummy class if praxis is unavailable."""
+
+    class BaseModel:
+        """Dummy class if praxis is unavailable."""
+
+
+PAD_VAL = 1123581321.0
+DEFAULT_QUANTILES = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
+
+# NestedMap keys
+_INPUT_TS = "input_ts"
+_TARGET_FUTURE = "actual_ts"
+_INPUT_PADDING = "input_padding"
+_OUTPUT_TS = "output_ts"
+_FREQ = "freq"
+_OUTPUT_TOKENS = "output_tokens"
+_STATS = "stats"
+
+
+# Small numerical value.
+_TOLERANCE = 1e-7
+
+
+def _shift_padded_seq(mask, seq):
+    num = seq.shape[1]
+
+    # Find the index of the first 0 in each row of the mask
+    first_zero_idx = jnp.argmin(mask, axis=1)
+
+    # Create a range array for indexing
+    idx_range = jnp.arange(num)
+
+    def shift_row(carry, x):
+        seq_row, shift = x
+        shifted_idx = (idx_range - shift) % num
+        shifted_row = seq_row[shifted_idx]
+        return carry, shifted_row
+
+    # Use lax.scan to shift each row of seq based on the corresponding
+    # first_zero_idx.
+    _, shifted_seq = lax.scan(shift_row, None, (seq, first_zero_idx))
+
+    return shifted_seq
+
+
+class ResidualBlock(BaseLayer):
+    """ResidualBlock."""
+
+    if _check_soft_dependencies("praxis", severity="none"):
+        # These should not be changed or removed
+        input_dims: int = 0
+        hidden_dims: int = 0
+        output_dims: int = 0
+        dropout_prob: float = 0.0
+        layer_norm: bool = False
+        dropout_tpl: LayerTpl = template_field(stochastics.Dropout)
+        ln_tpl: LayerTpl = template_field(normalizations.LayerNorm)
+        act_tpl: LayerTpl = template_field(activations.Swish)
+
+    def setup(self):
+        """Set up."""
+        lnorm_tpl = self.ln_tpl.clone()
+        lnorm_tpl.dim = self.output_dims
+        self.create_child("ln_layer", lnorm_tpl)
+
+        dropout_tpl = self.dropout_tpl.clone()
+        dropout_tpl.keep_prob = 1.0 - self.dropout_prob
+        self.create_child("dropout", dropout_tpl)
+
+        self.create_child(
+            "hidden_layer",
+            pax_fiddle.Config(
+                linears.FeedForward,
+                input_dims=self.input_dims,
+                output_dims=self.hidden_dims,
+                activation_tpl=self.act_tpl.clone(),
+            ),
+        )
+
+        self.create_child(
+            "output_layer",
+            pax_fiddle.Config(
+                linears.FeedForward,
+                input_dims=self.hidden_dims,
+                output_dims=self.output_dims,
+                activation_tpl=pax_fiddle.Config(activations.Identity),
+            ),
+        )
+
+        self.create_child(
+            "residual_layer",
+            pax_fiddle.Config(
+                linears.FeedForward,
+                input_dims=self.input_dims,
+                output_dims=self.output_dims,
+                activation_tpl=pax_fiddle.Config(activations.Identity),
+            ),
+        )
+
+    def __call__(self, inputs):
+        """__call__."""
+        hidden = self.hidden_layer(inputs)
+        output = self.output_layer(hidden)
+        output = self.dropout(output)
+        residual = self.residual_layer(inputs)
+        if self.layer_norm:
+            return self.ln_layer(output + residual)
+        else:
+            return output + residual
+
+
+def _masked_mean_std(inputs, padding):
+    pad_sum = jnp.sum(1 - padding, axis=2)
+
+    def _get_patch_index(arr):
+        indices = jnp.argmax(arr >= 3, axis=1)
+        row_sum = (arr >= 3).sum(axis=1)
+        return jnp.where(row_sum == 0, arr.shape[1] - 1, indices)
+
+    patch_indices = _get_patch_index(pad_sum)
+    bidxs = jnp.arange(inputs.shape[0])
+
+    arr = inputs[bidxs, patch_indices, :]
+    pad = padding[bidxs, patch_indices, :]
+
+    # Create a mask where P is 0
+    mask = 1 - pad
+
+    # Calculate the number of valid elements
+    num_valid_elements = jnp.sum(mask, axis=1)
+
+    num_valid_elements = jnp.where(num_valid_elements == 0, 1, num_valid_elements)
+
+    # Calculate the masked sum and squared sum of M
+    masked_sum = jnp.sum(arr * mask, axis=1)
+    masked_squared_sum = jnp.sum((arr * mask) ** 2, axis=1)
+
+    # Calculate the masked mean and standard deviation
+    masked_mean = masked_sum / num_valid_elements
+    masked_var = masked_squared_sum / num_valid_elements - masked_mean**2
+    masked_var = jnp.where(masked_var < 0.0, 0.0, masked_var)
+    masked_std = jnp.sqrt(masked_var)
+
+    return masked_mean, masked_std
+
+
+def _create_quantiles():
+    return DEFAULT_QUANTILES
+
+
+class PatchedTimeSeriesDecoder(BaseLayer):
+    """PatchedTimeSeriesDecoder."""
+
+    if _check_soft_dependencies("praxis", severity="none"):
+        # These should not be changed or removed
+        patch_len: int = 0
+        horizon_len: int = 0
+        model_dims: int = 0
+        hidden_dims: int = 0
+        quantiles: list[float] = dataclasses.field(default_factory=_create_quantiles)
+        residual_block_tpl: LayerTpl = template_field(ResidualBlock)
+        stacked_transformer_params_tpl: LayerTpl = template_field(
+            transformers.StackedTransformer
+        )
+        use_freq: bool = True
+
+    def setup(self):
+        """Construct the model."""
+        num_outputs = len(self.quantiles) + 1
+
+        stl = self.stacked_transformer_params_tpl.clone()
+        stl.model_dims = self.model_dims
+        stl.hidden_dims = self.hidden_dims
+        stl.mask_self_attention = True
+
+        self.create_child("stacked_transformer_layer", stl)
+
+        input_resl = self.residual_block_tpl.clone()
+        ff_in_dims = 2 * self.patch_len
+        input_resl.input_dims = ff_in_dims
+        input_resl.hidden_dims = self.hidden_dims
+        input_resl.output_dims = self.model_dims
+        self.create_child(
+            "input_ff_layer",
+            input_resl,
+        )
+
+        horizon_resl = self.residual_block_tpl.clone()
+        horizon_resl.input_dims = self.model_dims
+        horizon_resl.hidden_dims = self.hidden_dims
+        horizon_resl.output_dims = self.horizon_len * num_outputs
+        self.create_child(
+            "horizon_ff_layer",
+            horizon_resl,
+        )
+
+        self.create_child(
+            "position_emb",
+            pax_fiddle.Config(
+                layers.PositionalEmbedding, embedding_dims=self.model_dims
+            ),
+        )
+
+        if self.use_freq:
+            self.create_child(
+                "freq_emb",
+                pax_fiddle.Config(
+                    embedding_softmax.Embedding,
+                    num_classes=3,
+                    input_dims=self.model_dims,
+                ),
+            )
+
+    def transform_decode_state(self, transform_fn):
+        """transform_decode_state."""
+        self.stacked_transformer_layer.transform_decode_state(transform_fn)
+
+    def _forward_transform(self, inputs, patched_pads):
+        """Input is of shape [B, N, P]."""
+        mu, sigma = _masked_mean_std(inputs, patched_pads)
+        sigma = jnp.where(sigma < _TOLERANCE, 1.0, sigma)
+        # Normalize each patch.
+        outputs = (inputs - mu[:, None, None]) / sigma[:, None, None]
+        outputs = jnp.where(jnp.abs(inputs - PAD_VAL) < _TOLERANCE, PAD_VAL, outputs)
+        return outputs, (mu, sigma)
+
+    def _reverse_transform(self, outputs, stats):
+        """Output is of shape [B, N, P, Q]."""
+        mu, sigma = stats
+        return outputs * sigma[:, None, None, None] + mu[:, None, None, None]
+
+    def _preprocess_input(
+        self,
+        input_ts,
+        input_padding,
+        pos_emb=None,
+    ):
+        """Preprocess input for stacked transformer."""
+        # Reshape into patches.
+        patched_inputs = es.jax_einshape("b(np)->bnp", input_ts, p=self.patch_len)
+        patched_pads = es.jax_einshape("b(np)->bnp", input_padding, p=self.patch_len)
+        patched_inputs = jnp.where(
+            jnp.abs(patched_pads - 1.0) < _TOLERANCE, 0.0, patched_inputs
+        )
+        patched_pads = jnp.where(
+            jnp.abs(patched_inputs - PAD_VAL) < _TOLERANCE, 1, patched_pads
+        )
+        patched_inputs, stats = self._forward_transform(patched_inputs, patched_pads)
+
+        # B x N x D
+        patched_inputs = patched_inputs * (1.0 - patched_pads)
+        concat_inputs = jnp.concatenate([patched_inputs, patched_pads], axis=-1)
+        model_input = self.input_ff_layer(concat_inputs)
+        # A patch should not be padded even if there is at least one zero.
+        patched_padding = jnp.min(patched_pads, axis=-1)
+
+        if pos_emb is None:
+            position_emb = self.position_emb(seq_length=model_input.shape[1])
+        else:
+            position_emb = pos_emb
+        if self.do_eval:
+            if position_emb.shape[0] != model_input.shape[0]:
+                position_emb = jnp.repeat(position_emb, model_input.shape[0], axis=0)
+            position_emb = _shift_padded_seq(patched_padding, position_emb)
+        model_input += position_emb
+
+        return model_input, patched_padding, stats, patched_inputs
+
+    def _postprocess_output(
+        self,
+        model_output,
+        num_outputs,
+        stats,
+    ):
+        """Postprocess output of stacked transformer."""
+        # B x N x (H.Q)
+        output_ts = self.horizon_ff_layer(model_output)
+        output_ts = es.jax_einshape(
+            "bn(hq)->bnhq", output_ts, q=num_outputs, h=self.horizon_len
+        )
+        return self._reverse_transform(output_ts, stats)
+
+    def __call__(self, inputs):
+        """PatchTST call.
+
+        Args:
+          inputs: A NestedMap containing (1) input_ts: input sequence of shape [B,
+            T] where T must be multiple of patch_length; (2) input_padding: that
+            contains padding map.
+
+        Returns
+        -------
+          A nested map with two keys:
+          (1) 'output_tokens' of shape [B, N, D].
+          (2) 'output_ts' of shape [B, N, H, Q]
+          (3) 'stats' a tuple of statistics for renormalization.
+        """
+        input_ts, input_padding = inputs[_INPUT_TS], inputs[_INPUT_PADDING]
+        num_outputs = len(self.quantiles) + 1
+        model_input, patched_padding, stats, _ = self._preprocess_input(
+            input_ts=input_ts,
+            input_padding=input_padding,
+        )
+        if self.use_freq:
+            freq = inputs[_FREQ].astype(jnp.int32)
+            f_emb = self.freq_emb(freq)  # B x 1 x D
+            f_emb = jnp.repeat(f_emb, model_input.shape[1], axis=1)
+            model_input += f_emb
+        model_output = self.stacked_transformer_layer(model_input, patched_padding)
+
+        output_ts = self._postprocess_output(model_output, num_outputs, stats)
+        return NestedMap(
+            {_OUTPUT_TOKENS: model_output, _OUTPUT_TS: output_ts, _STATS: stats}
+        )
+
+    def decode(
+        self,
+        inputs,
+        horizon_len,
+        output_patch_len=None,
+        max_len=512,
+        return_forecast_on_context=False,
+    ):
+        """Decode."""
+        final_out = inputs[_INPUT_TS]
+        context_len = final_out.shape[1]
+        paddings = inputs[_INPUT_PADDING]
+        if self.use_freq:
+            freq = inputs[_FREQ].astype(jnp.int32)
+        else:
+            freq = jnp.zeros([final_out.shape[0], 1], dtype=jnp.int32)
+        full_outputs = []
+        if paddings.shape[1] != final_out.shape[1] + horizon_len:
+            raise ValueError(
+                "Length of paddings must match length of input + horizon_len:"
+                f" {paddings.shape[1]} != {final_out.shape[1]} + {horizon_len}"
+            )
+        if output_patch_len is None:
+            output_patch_len = self.horizon_len
+        num_decode_patches = (horizon_len + output_patch_len - 1) // output_patch_len
+        for step_index in range(num_decode_patches):
+            current_padding = paddings[:, 0 : final_out.shape[1]]
+            input_ts = final_out[:, -max_len:]
+            input_padding = current_padding[:, -max_len:]
+            model_input = NestedMap(
+                input_ts=input_ts,
+                input_padding=input_padding,
+                freq=freq,
+            )
+            fprop_outputs = self(model_input)[_OUTPUT_TS]
+            if return_forecast_on_context and step_index == 0:
+                # For the first decodings step, collect the model forecast on the
+                # context except the unavailable first input batch forecast.
+                new_full_ts = fprop_outputs[:, :-1, : self.patch_len, :]
+                new_full_ts = es.jax_einshape("bnph->b(np)h", new_full_ts)
+
+                full_outputs.append(new_full_ts)
+
+            # (full batch, last patch, output_patch_len, index of mean forecast = 0)
+            new_ts = fprop_outputs[:, -1, :output_patch_len, 0]
+            new_full_ts = fprop_outputs[:, -1, :output_patch_len, :]
+            # (full batch, last patch, output_patch_len, all output indices)
+            full_outputs.append(new_full_ts)
+            final_out = jnp.concatenate([final_out, new_ts], axis=-1)
+
+        if return_forecast_on_context:
+            # `full_outputs` indexing starts at after the first input patch.
+            full_outputs = jnp.concatenate(full_outputs, axis=1)[
+                :, : (context_len - self.patch_len + horizon_len), :
+            ]
+        else:
+            # `full_outputs` indexing starts at the forecast horizon.
+            full_outputs = jnp.concatenate(full_outputs, axis=1)[:, 0:horizon_len, :]
+
+        return (full_outputs[:, :, 0], full_outputs)
+
+
+class PatchedDecoderFinetuneModel(BaseModel):
+    """PatchedDecoderFinetuneModel."""
+
+    def __init__(self):
+        self.core_layer_tpl = template_field(PatchedTimeSeriesDecoder)
+        self.freq: int = 0
+
+    def setup(self):
+        """Set up."""
+        self.create_child("core_layer", self.core_layer_tpl)
+
+    def compute_predictions(self, input_batch):
+        """compute_predictions."""
+        input_ts = input_batch[_INPUT_TS]
+        input_padding = jnp.zeros_like(input_ts)
+        context_len = input_ts.shape[1]
+        input_patch_len = self.core_layer_tpl.patch_len
+        context_pad = (
+            (context_len + input_patch_len - 1) // input_patch_len
+        ) * input_patch_len - context_len
+
+        input_ts = jnp.pad(input_ts, [(0, 0), (context_pad, 0)])
+        input_padding = jnp.pad(
+            input_padding, [(0, 0), (context_pad, 0)], constant_values=1
+        )
+        freq = jnp.ones([input_ts.shape[0], 1], dtype=jnp.int32) * self.freq
+        new_input_batch = NestedMap(
+            input_ts=input_ts,
+            input_padding=input_padding,
+            freq=freq,
+        )
+        return self.core_layer(new_input_batch)
+
+    def _quantile_loss(self, pred, actual, quantile: float):
+        dev = actual - pred
+        loss_first = dev * quantile
+        loss_second = -dev * (1.0 - quantile)
+        return 2 * jnp.where(loss_first >= 0, loss_first, loss_second)
+
+    def compute_loss(self, prediction_output, input_batch):
+        """compute_loss."""
+        output_ts = prediction_output[_OUTPUT_TS]
+        actual_ts = input_batch[_TARGET_FUTURE]
+        pred_ts = output_ts[:, -1, 0 : actual_ts.shape[1], :]
+        loss = jnp.square(pred_ts[:, :, 0] - actual_ts)
+        for i, quantile in enumerate(self.core_layer.quantiles):
+            loss += self._quantile_loss(pred_ts[:, :, i + 1], actual_ts, quantile)
+        loss = loss.mean()
+        loss_weight = jnp.array(1.0, dtype=jnp.float32)
+        per_example_out = NestedMap()
+        return {"avg_qloss": (loss, loss_weight)}, per_example_out
diff --git a/sktime/libs/timesfm/time_features.py b/sktime/libs/timesfm/time_features.py
new file mode 100644
index 000000000..7843ab7e1
--- /dev/null
+++ b/sktime/libs/timesfm/time_features.py
@@ -0,0 +1,209 @@
+# Copyright 2024 The Google Research Authors.
+#
+# 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.
+
+"""time_features."""
+
+import numpy as np
+import pandas as pd
+from pandas.tseries.holiday import (
+    SU,
+    TH,
+    EasterMonday,
+    GoodFriday,
+    Holiday,
+    USColumbusDay,
+    USLaborDay,
+    USMartinLutherKingJr,
+    USMemorialDay,
+    USPresidentsDay,
+    USThanksgivingDay,
+)
+from pandas.tseries.offsets import DateOffset, Day, Easter
+from sklearn.preprocessing import StandardScaler
+
+# avoid another dependency
+# from tqdm import tqdm
+
+# This is 183 to cover half a year (in both directions), also for leap years
+# + 17 as Eastern can be between March, 22 - April, 25
+MAX_WINDOW = 183 + 17
+
+
+def _distance_to_holiday(holiday):
+    """Return distance to given holiday."""
+
+    def _distance_to_day(index):
+        holiday_date = holiday.dates(
+            index - pd.Timedelta(days=MAX_WINDOW),
+            index + pd.Timedelta(days=MAX_WINDOW),
+        )
+        assert (
+            len(holiday_date) != 0  # pylint: disable=g-explicit-length-test
+        ), f"No closest holiday for the date index {index} found."
+        # It sometimes returns two dates if it is exactly half a year after the
+        # holiday. In this case, the smaller distance (182 days) is returned.
+        return (index - holiday_date[0]).days
+
+    return _distance_to_day
+
+
+EasterSunday = Holiday("Easter Sunday", month=1, day=1, offset=[Easter(), Day(0)])
+NewYearsDay = Holiday("New Years Day", month=1, day=1)
+SuperBowl = Holiday("Superbowl", month=2, day=1, offset=DateOffset(weekday=SU(1)))
+MothersDay = Holiday("Mothers Day", month=5, day=1, offset=DateOffset(weekday=SU(2)))
+IndependenceDay = Holiday("Independence Day", month=7, day=4)
+ChristmasEve = Holiday("Christmas", month=12, day=24)
+ChristmasDay = Holiday("Christmas", month=12, day=25)
+NewYearsEve = Holiday("New Years Eve", month=12, day=31)
+BlackFriday = Holiday(
+    "Black Friday",
+    month=11,
+    day=1,
+    offset=[pd.DateOffset(weekday=TH(4)), Day(1)],
+)
+CyberMonday = Holiday(
+    "Cyber Monday",
+    month=11,
+    day=1,
+    offset=[pd.DateOffset(weekday=TH(4)), Day(4)],
+)
+
+HOLIDAYS = [
+    EasterMonday,
+    GoodFriday,
+    USColumbusDay,
+    USLaborDay,
+    USMartinLutherKingJr,
+    USMemorialDay,
+    USPresidentsDay,
+    USThanksgivingDay,
+    EasterSunday,
+    NewYearsDay,
+    SuperBowl,
+    MothersDay,
+    IndependenceDay,
+    ChristmasEve,
+    ChristmasDay,
+    NewYearsEve,
+    BlackFriday,
+    CyberMonday,
+]
+
+
+class TimeCovariates:
+    """TimeCovariates."""
+
+    def __init__(
+        self,
+        datetimes,
+        normalized=True,
+        holiday=False,
+    ):
+        """Init function.
+
+        Args:
+          datetimes: pandas DatetimeIndex (lowest granularity supported is min)
+          normalized: whether to normalize features or not
+          holiday: fetch holiday features or not
+
+        Returns
+        -------
+          None
+        """
+        self.normalized = normalized
+        self.dti = datetimes
+        self.holiday = holiday
+
+    def _minute_of_hour(self):
+        minutes = np.array(self.dti.minute, dtype=np.float32)
+        if self.normalized:
+            minutes = minutes / 59.0 - 0.5
+        return minutes
+
+    def _hour_of_day(self):
+        hours = np.array(self.dti.hour, dtype=np.float32)
+        if self.normalized:
+            hours = hours / 23.0 - 0.5
+        return hours
+
+    def _day_of_week(self):
+        day_week = np.array(self.dti.dayofweek, dtype=np.float32)
+        if self.normalized:
+            day_week = day_week / 6.0 - 0.5
+        return day_week
+
+    def _day_of_month(self):
+        day_month = np.array(self.dti.day, dtype=np.float32)
+        if self.normalized:
+            day_month = day_month / 30.0 - 0.5
+        return day_month
+
+    def _day_of_year(self):
+        day_year = np.array(self.dti.dayofyear, dtype=np.float32)
+        if self.normalized:
+            day_year = day_year / 364.0 - 0.5
+        return day_year
+
+    def _month_of_year(self):
+        month_year = np.array(self.dti.month, dtype=np.float32)
+        if self.normalized:
+            month_year = month_year / 11.0 - 0.5
+        return month_year
+
+    def _week_of_year(self):
+        week_year = np.array(self.dti.strftime("%U").astype(int), dtype=np.float32)
+        if self.normalized:
+            week_year = week_year / 51.0 - 0.5
+        return week_year
+
+    def _get_holidays(self):
+        dti_series = self.dti.to_series()
+        hol_variates = np.vstack(
+            # [dti_series.apply(_distance_to_holiday(h)).values for h in tqdm(HOLIDAYS)]
+            [dti_series.apply(_distance_to_holiday(h)).values for h in HOLIDAYS]
+        )
+        # hol_variates is (num_holiday, num_time_steps), the normalization should be
+        # performed in the num_time_steps dimension.
+        return StandardScaler().fit_transform(hol_variates.T).T
+
+    def get_covariates(self):
+        """Get all time covariates."""
+        moh = self._minute_of_hour().reshape(1, -1)
+        hod = self._hour_of_day().reshape(1, -1)
+        dom = self._day_of_month().reshape(1, -1)
+        dow = self._day_of_week().reshape(1, -1)
+        doy = self._day_of_year().reshape(1, -1)
+        moy = self._month_of_year().reshape(1, -1)
+        woy = self._week_of_year().reshape(1, -1)
+
+        all_covs = [
+            moh,
+            hod,
+            dom,
+            dow,
+            doy,
+            moy,
+            woy,
+        ]
+        columns = ["moh", "hod", "dom", "dow", "doy", "moy", "woy"]
+        if self.holiday:
+            hol_covs = self._get_holidays()
+            all_covs.append(hol_covs)
+            columns += [f"hol_{i}" for i in range(len(HOLIDAYS))]
+
+        return pd.DataFrame(
+            data=np.vstack(all_covs).transpose(),
+            columns=columns,
+            index=self.dti,
+        )
diff --git a/sktime/libs/timesfm/timesfm.py b/sktime/libs/timesfm/timesfm.py
new file mode 100644
index 000000000..4d5d35e72
--- /dev/null
+++ b/sktime/libs/timesfm/timesfm.py
@@ -0,0 +1,722 @@
+# Copyright 2024 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.
+
+"""TimesFM forecast API for inference."""
+
+import collections
+import logging
+import multiprocessing
+import time
+from os import path
+
+from sktime.utils.dependencies import _check_soft_dependencies
+
+if _check_soft_dependencies("einshape", severity="none"):
+    import einshape as es
+
+if _check_soft_dependencies("jax", severity="none"):
+    import jax
+    import jax.numpy as jnp
+
+if _check_soft_dependencies("huggingface-hub", severity="none"):
+    from huggingface_hub import snapshot_download
+
+if _check_soft_dependencies("paxml", severity="none"):
+    from paxml import checkpoints, tasks_lib
+
+    FLAX = checkpoints.CheckpointType.FLAX
+else:
+    FLAX = None
+
+if _check_soft_dependencies("praxis", severity="none"):
+    from praxis import base_hyperparams, base_layer, pax_fiddle, py_utils, pytypes
+    from praxis.layers import normalizations, transformers
+
+    instantiate = base_hyperparams.instantiate
+    NestedMap = py_utils.NestedMap
+    JTensor = pytypes.JTensor
+
+
+if _check_soft_dependencies("utilsforecast", severity="none"):
+    from utilsforecast.processing import make_future_dataframe
+
+
+import numpy as np
+
+from sktime.libs.timesfm import patched_decoder, xreg_lib
+
+_TOL = 1e-6
+
+
+def process_group(key, group, value_name, forecast_context_len):
+    """process_group."""
+    group = group.tail(forecast_context_len)
+    return np.array(group[value_name], dtype=np.float32), key
+
+
+def moving_average(arr, window_size):
+    """moving_average."""
+    # Pad with zeros to handle initial window positions
+    arr_padded = np.pad(arr, (window_size - 1, 0), "constant")
+    smoothed_arr = np.convolve(arr_padded, np.ones(window_size), "valid") / window_size
+    return [smoothed_arr, arr - smoothed_arr]
+
+
+def freq_map(freq):
+    """freq_map."""
+    freq = str.upper(freq)
+    if (
+        freq.endswith("H")
+        or freq.endswith("T")
+        or freq.endswith("MIN")
+        or freq.endswith("D")
+        or freq.endswith("B")
+        or freq.endswith("U")
+    ):
+        return 0
+    elif freq.endswith(("W", "M", "MS")):
+        return 1
+    elif freq.endswith("Y") or freq.endswith("Q"):
+        return 2
+    else:
+        raise ValueError(f"Invalid frequency: {freq}")
+
+
+# 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 TimesFm:
+    """TimesFM."""
+
+    def _logging(self, s):
+        if self._verbose:
+            print(s)
+
+    def __init__(
+        self,
+        context_len,
+        horizon_len,
+        input_patch_len,
+        output_patch_len,
+        num_layers,
+        model_dims,
+        per_core_batch_size=32,
+        backend="cpu",
+        quantiles=None,
+        verbose=True,
+    ):
+        self.per_core_batch_size = per_core_batch_size
+        self.backend = backend
+        self.num_devices = jax.local_device_count(self.backend)
+        self.global_batch_size = self.per_core_batch_size * self.num_devices
+
+        self.context_len = context_len
+        self.horizon_len = horizon_len
+        self.input_patch_len = input_patch_len
+        self.output_patch_len = output_patch_len
+        self._horizon_start = self.context_len - self.input_patch_len
+
+        self.mesh_shape = [1, self.num_devices, 1]
+        self.mesh_name = ["replica", "data", "mdl"]
+        if quantiles is None:
+            quantiles = patched_decoder.DEFAULT_QUANTILES
+
+        self.model_p = pax_fiddle.Config(
+            patched_decoder.PatchedTimeSeriesDecoder,
+            name="patched_decoder",
+            horizon_len=self.output_patch_len,
+            patch_len=input_patch_len,
+            model_dims=model_dims,
+            hidden_dims=model_dims,
+            residual_block_tpl=pax_fiddle.Config(patched_decoder.ResidualBlock),
+            quantiles=quantiles,
+            use_freq=True,
+            stacked_transformer_params_tpl=pax_fiddle.Config(
+                transformers.StackedTransformer,
+                num_heads=16,
+                num_layers=num_layers,
+                transformer_layer_params_tpl=pax_fiddle.Config(
+                    transformers.Transformer,
+                    ln_tpl=pax_fiddle.Config(
+                        normalizations.RmsNorm,
+                    ),
+                ),
+            ),
+        )
+
+        self._key1, self._key2 = jax.random.split(jax.random.PRNGKey(42))
+        self._model = None
+        self._train_state = None
+        self._pmapped_decode = None
+        self._verbose = verbose
+        self._eval_context = base_layer.JaxContext.HParams(do_eval=True)
+        try:
+            multiprocessing.set_start_method("spawn")
+        except RuntimeError:
+            print("Multiprocessing context has already been set.")
+
+    def _get_sample_inputs(self):
+        return {
+            "input_ts": jnp.zeros(
+                (
+                    self.per_core_batch_size,
+                    self.context_len + self.output_patch_len,
+                ),
+                dtype=jnp.float32,
+            ),
+            "input_padding": jnp.zeros(
+                (
+                    self.per_core_batch_size,
+                    self.context_len + self.output_patch_len,
+                ),
+                dtype=jnp.float32,
+            ),
+            "freq": jnp.zeros(
+                (
+                    self.per_core_batch_size,
+                    1,
+                ),
+                dtype=jnp.int32,
+            ),
+        }
+
+    def load_from_checkpoint(
+        self,
+        checkpoint_path=None,
+        repo_id="google/timesfm-1.0-200m",
+        checkpoint_type=FLAX,
+        step=None,
+    ):
+        """load_from_checkpoint."""
+        # Download the checkpoint from Hugging Face Hub if not given
+        if checkpoint_path is None:
+            checkpoint_path = path.join(snapshot_download(repo_id), "checkpoints")
+
+        #  Initialize the model weights.
+        self._logging("Constructing model weights.")
+        start_time = time.time()
+        self._model = instantiate(self.model_p)
+        var_weight_hparams = self._model.abstract_init_with_metadata(
+            self._get_sample_inputs(), do_eval=True
+        )
+        train_state_partition_specs = tasks_lib.create_state_partition_specs(
+            var_weight_hparams,
+            mesh_shape=self.mesh_shape,
+            mesh_axis_names=self.mesh_name,
+            discard_opt_states=True,
+            learners=None,
+        )
+        train_state_local_shapes = tasks_lib.create_state_unpadded_shapes(
+            var_weight_hparams,
+            discard_opt_states=True,
+            learners=None,
+        )
+        self._logging(
+            f"Constructed model weights in {time.time() - start_time:.2f} seconds."
+        )
+
+        # Load the model weights.
+        self._logging(f"Restoring checkpoint from {checkpoint_path}.")
+        start_time = time.time()
+        self._train_state = checkpoints.restore_checkpoint(
+            train_state_local_shapes,
+            checkpoint_dir=checkpoint_path,
+            checkpoint_type=checkpoint_type,
+            state_specs=train_state_partition_specs,
+            step=step,
+        )
+        self._logging(f"Restored checkpoint in {time.time() - start_time:.2f} seconds.")
+        self.jit_decode()
+
+    def jit_decode(self):
+        """Jitting decoding function."""
+
+        # Initialize and jit the decode fn.
+        def _decode(inputs):
+            assert self._model is not None
+            assert self._train_state is not None
+            return self._model.apply(
+                self._train_state.mdl_vars,
+                inputs,
+                horizon_len=self.horizon_len,
+                output_patch_len=self.output_patch_len,
+                max_len=self.context_len,
+                return_forecast_on_context=True,
+                rngs={
+                    base_layer.PARAMS: self._key1,
+                    base_layer.RANDOM: self._key2,
+                },
+                method=self._model.decode,
+            )
+
+        self._logging("Jitting decoding.")
+        start_time = time.time()
+        self._pmapped_decode = jax.pmap(
+            _decode,
+            axis_name="batch",
+            devices=jax.devices(self.backend),
+            backend=self.backend,
+            axis_size=self.num_devices,
+        )
+        with base_layer.JaxContext.new_context(hparams=self._eval_context):
+            _ = self._pmapped_decode(
+                NestedMap(
+                    {
+                        "input_ts": jnp.zeros(
+                            (
+                                self.num_devices,
+                                self.per_core_batch_size,
+                                self.context_len,
+                            ),
+                            dtype=jnp.float32,
+                        ),
+                        "input_padding": jnp.zeros(
+                            (
+                                self.num_devices,
+                                self.per_core_batch_size,
+                                self.context_len + self.horizon_len,
+                            ),
+                            dtype=jnp.float32,
+                        ),
+                        "date_features": None,
+                        "freq": jnp.zeros(
+                            (self.num_devices, self.per_core_batch_size, 1),
+                            dtype=jnp.int32,
+                        ),
+                    }
+                )
+            )
+        self._logging(f"Jitted decoding in {time.time() - start_time:.2f} seconds.")
+
+    def _preprocess(self, inputs, freq):
+        input_ts, input_padding, inp_freq = [], [], []
+
+        pmap_pad = (
+            (len(inputs) - 1) // self.global_batch_size + 1
+        ) * self.global_batch_size - len(inputs)
+
+        for i, ts in enumerate(inputs):
+            input_len = ts.shape[0]
+            padding = np.zeros(shape=(input_len + self.horizon_len,), dtype=float)
+            if input_len < self.context_len:
+                num_front_pad = self.context_len - input_len
+                ts = np.concatenate(
+                    [np.zeros(shape=(num_front_pad,), dtype=float), ts], axis=0
+                )
+                padding = np.concatenate(
+                    [np.ones(shape=(num_front_pad,), dtype=float), padding], axis=0
+                )
+            elif input_len > self.context_len:
+                ts = ts[-self.context_len :]
+                padding = padding[-(self.context_len + self.horizon_len) :]
+
+            input_ts.append(ts)
+            input_padding.append(padding)
+            inp_freq.append(freq[i])
+
+        # Padding the remainder batch.
+        for _ in range(pmap_pad):
+            input_ts.append(input_ts[-1])
+            input_padding.append(input_padding[-1])
+            inp_freq.append(inp_freq[-1])
+
+        return (
+            np.stack(input_ts, axis=0),
+            np.stack(input_padding, axis=0),
+            np.array(inp_freq).astype(np.int32).reshape(-1, 1),
+            pmap_pad,
+        )
+
+    def forecast(
+        self,
+        inputs,
+        freq=None,
+        window_size=None,
+        forecast_context_len=None,
+        return_forecast_on_context=False,
+    ):
+        """Forecast."""
+        if not self._train_state or not self._model:
+            raise ValueError(
+                "Checkpoint not loaded. Call `load_from_checkpoint` before"
+                " `forecast`."
+            )
+        if forecast_context_len is None:
+            forecast_context_len = self.context_len
+        inputs = [np.array(ts)[-forecast_context_len:] for ts in inputs]
+        inp_min = np.min([np.min(ts) for ts in inputs])
+
+        if window_size is not None:
+            new_inputs = []
+            for ts in inputs:
+                new_inputs.extend(moving_average(ts, window_size))
+            inputs = new_inputs
+
+        if freq is None:
+            logging.info("No frequency provided via `freq`. Default to high (0).")
+            freq = [0] * len(inputs)
+
+        input_ts, input_padding, inp_freq, pmap_pad = self._preprocess(inputs, freq)
+        with base_layer.JaxContext.new_context(hparams=self._eval_context):
+            mean_outputs = []
+            full_outputs = []
+            assert input_ts.shape[0] % self.global_batch_size == 0
+            for i in range(input_ts.shape[0] // self.global_batch_size):
+                input_ts_in = jnp.array(
+                    input_ts[
+                        i * self.global_batch_size : (i + 1) * self.global_batch_size
+                    ]
+                )
+                input_padding_in = jnp.array(
+                    input_padding[
+                        i * self.global_batch_size : (i + 1) * self.global_batch_size
+                    ],
+                )
+                inp_freq_in = jnp.array(
+                    inp_freq[
+                        i * self.global_batch_size : (i + 1) * self.global_batch_size, :
+                    ],
+                    dtype=jnp.int32,
+                )
+                pmapped_inputs = NestedMap(
+                    {
+                        "input_ts": es.jax_einshape(
+                            "(db)...->db...",
+                            input_ts_in,
+                            d=self.num_devices,
+                        ),
+                        "input_padding": es.jax_einshape(
+                            "(db)...->db...",
+                            input_padding_in,
+                            d=self.num_devices,
+                        ),
+                        "date_features": None,
+                        "freq": es.jax_einshape(
+                            "(db)...->db...",
+                            inp_freq_in,
+                            d=self.num_devices,
+                        ),
+                    }
+                )
+                mean_output, full_output = self._pmapped_decode(pmapped_inputs)
+                if not return_forecast_on_context:
+                    mean_output = mean_output[:, :, self._horizon_start :, ...]
+                    full_output = full_output[:, :, self._horizon_start :, ...]
+                mean_output = es.jax_einshape(
+                    "db...->(db)...", mean_output, d=self.num_devices
+                )
+                full_output = es.jax_einshape(
+                    "db...->(db)...", full_output, d=self.num_devices
+                )
+                mean_output = np.array(mean_output)
+                full_output = np.array(full_output)
+                mean_outputs.append(mean_output)
+                full_outputs.append(full_output)
+
+        mean_outputs = np.concatenate(mean_outputs, axis=0)
+        full_outputs = np.concatenate(full_outputs, axis=0)
+
+        if pmap_pad > 0:
+            mean_outputs = mean_outputs[:-pmap_pad, ...]
+            full_outputs = full_outputs[:-pmap_pad, ...]
+
+        if window_size is not None:
+            mean_outputs = mean_outputs[0::2, ...] + mean_outputs[1::2, ...]
+            full_outputs = full_outputs[0::2, ...] + full_outputs[1::2, ...]
+        if inp_min >= 0:
+            mean_outputs = np.maximum(mean_outputs, 0.0)
+            full_outputs = np.maximum(full_outputs, 0.0)
+        return mean_outputs, full_outputs
+
+    def forecast_with_covariates(
+        self,
+        inputs,
+        dynamic_numerical_covariates=None,
+        dynamic_categorical_covariates=None,
+        static_numerical_covariates=None,
+        static_categorical_covariates=None,
+        freq=None,
+        window_size=None,
+        forecast_context_len=None,
+        xreg_mode="xreg + timesfm",
+        normalize_xreg_target_per_input=True,
+        ridge=0.0,
+        max_rows_per_col=0,
+        force_on_cpu=False,
+    ):
+        """forecast_with_covariates."""
+        # Verify and bookkeep covariates.
+        if not (
+            dynamic_numerical_covariates
+            or dynamic_categorical_covariates
+            or static_numerical_covariates
+            or static_categorical_covariates
+        ):
+            raise ValueError(
+                "At least one of dynamic_numerical_covariates,"
+                " dynamic_categorical_covariates, static_numerical_covariates,"
+                " static_categorical_covariates must be set."
+            )
+
+        # Track the lengths of (1) each input, (2) the part that can be used in the
+        # linear model, and (3) the horizon.
+        input_lens, train_lens, test_lens = [], [], []
+
+        for i, input_ts in enumerate(inputs):
+            input_len = len(input_ts)
+            input_lens.append(input_len)
+
+            if xreg_mode == "timesfm + xreg":
+                # For fitting residuals, no TimesFM forecast on the first patch.
+                train_lens.append(max(0, input_len - self.input_patch_len))
+            elif xreg_mode == "xreg + timesfm":
+                train_lens.append(input_len)
+            else:
+                raise ValueError(f"Unsupported mode: {xreg_mode}")
+
+            if dynamic_numerical_covariates:
+                test_lens.append(
+                    len(list(dynamic_numerical_covariates.values())[0][i]) - input_len
+                )
+            elif dynamic_categorical_covariates:
+                test_lens.append(
+                    len(list(dynamic_categorical_covariates.values())[0][i]) - input_len
+                )
+            else:
+                test_lens.append(self.horizon_len)
+
+            if test_lens[-1] > self.horizon_len:
+                raise ValueError(
+                    "Forecast requested longer horizon than the model definition "
+                    f"supports: {test_lens[-1]} vs {self.horizon_len}."
+                )
+
+        # Prepare the covariates into train and test.
+        train_dynamic_numerical_covariates = collections.defaultdict(list)
+        test_dynamic_numerical_covariates = collections.defaultdict(list)
+        train_dynamic_categorical_covariates = collections.defaultdict(list)
+        test_dynamic_categorical_covariates = collections.defaultdict(list)
+        for covariates, train_covariates, test_covariates in (
+            (
+                dynamic_numerical_covariates,
+                train_dynamic_numerical_covariates,
+                test_dynamic_numerical_covariates,
+            ),
+            (
+                dynamic_categorical_covariates,
+                train_dynamic_categorical_covariates,
+                test_dynamic_categorical_covariates,
+            ),
+        ):
+            if not covariates:
+                continue
+            for covariate_name, covariate_values in covariates.items():
+                for input_len, train_len, covariate_value in zip(
+                    input_lens, train_lens, covariate_values
+                ):
+                    train_covariates[covariate_name].append(
+                        covariate_value[(input_len - train_len) : input_len]
+                    )
+                    test_covariates[covariate_name].append(covariate_value[input_len:])
+
+        # Fit models.
+        if xreg_mode == "timesfm + xreg":
+            # Forecast via TimesFM then fit a model on the residuals.
+            mean_outputs, _ = self.forecast(
+                inputs,
+                freq,
+                window_size,
+                forecast_context_len,
+                return_forecast_on_context=True,
+            )
+            targets = [
+                (
+                    np.array(input_ts)[-train_len:]
+                    - mean_output[
+                        (self._horizon_start - train_len) : self._horizon_start
+                    ]
+                )
+                for input_ts, mean_output, train_len in zip(
+                    inputs, mean_outputs, train_lens
+                )
+            ]
+            per_instance_stats = None
+            if normalize_xreg_target_per_input:
+                targets, per_instance_stats = _normalize(targets)
+            xregs = xreg_lib.BatchedInContextXRegLinear(
+                targets=targets,
+                train_lens=train_lens,
+                test_lens=test_lens,
+                train_dynamic_numerical_covariates=train_dynamic_numerical_covariates,
+                test_dynamic_numerical_covariates=test_dynamic_numerical_covariates,
+                train_dynamic_categorical_covariates=train_dynamic_categorical_covariates,
+                test_dynamic_categorical_covariates=test_dynamic_categorical_covariates,
+                static_numerical_covariates=static_numerical_covariates,
+                static_categorical_covariates=static_categorical_covariates,
+            ).fit(
+                ridge=ridge,
+                one_hot_encoder_drop=None if ridge > 0 else "first",
+                max_rows_per_col=max_rows_per_col,
+                force_on_cpu=force_on_cpu,
+                debug_info=False,
+                assert_covariates=True,
+                assert_covariate_shapes=True,
+            )
+            if normalize_xreg_target_per_input:
+                xregs = _renormalize(xregs, per_instance_stats)
+            outputs = [
+                (
+                    mean_output[self._horizon_start : (self._horizon_start + test_len)]
+                    + xreg
+                )
+                for mean_output, test_len, xreg in zip(mean_outputs, test_lens, xregs)
+            ]
+
+        else:
+            # Fit a model on the targets then forecast on the residuals via TimesFM.
+            targets = [
+                np.array(input_ts)[-train_len:]
+                for input_ts, train_len in zip(inputs, train_lens)
+            ]
+            per_instance_stats = None
+            if normalize_xreg_target_per_input:
+                targets, per_instance_stats = _normalize(targets)
+            xregs, xregs_on_context, _, _, _ = xreg_lib.BatchedInContextXRegLinear(
+                targets=targets,
+                train_lens=train_lens,
+                test_lens=test_lens,
+                train_dynamic_numerical_covariates=train_dynamic_numerical_covariates,
+                test_dynamic_numerical_covariates=test_dynamic_numerical_covariates,
+                train_dynamic_categorical_covariates=train_dynamic_categorical_covariates,
+                test_dynamic_categorical_covariates=test_dynamic_categorical_covariates,
+                static_numerical_covariates=static_numerical_covariates,
+                static_categorical_covariates=static_categorical_covariates,
+            ).fit(
+                ridge=ridge,
+                one_hot_encoder_drop=None if ridge > 0 else "first",
+                max_rows_per_col=max_rows_per_col,
+                force_on_cpu=force_on_cpu,
+                debug_info=True,
+                assert_covariates=True,
+                assert_covariate_shapes=True,
+            )
+            mean_outputs, _ = self.forecast(
+                [
+                    target - xreg_on_context
+                    for target, xreg_on_context in zip(targets, xregs_on_context)
+                ],
+                freq,
+                window_size,
+                forecast_context_len,
+                return_forecast_on_context=True,
+            )
+            outputs = [
+                (
+                    mean_output[self._horizon_start : (self._horizon_start + test_len)]
+                    + xreg
+                )
+                for mean_output, test_len, xreg in zip(mean_outputs, test_lens, xregs)
+            ]
+            if normalize_xreg_target_per_input:
+                outputs = _renormalize(outputs, per_instance_stats)
+
+        return outputs, xregs
+
+    def forecast_on_df(
+        self,
+        inputs,
+        freq,
+        forecast_context_len=0,
+        value_name="values",
+        model_name="timesfm",
+        window_size=None,
+        num_jobs=1,
+        verbose=True,
+    ):
+        """forecast_on_df."""
+        if not (
+            "unique_id" in inputs.columns
+            and "ds" in inputs.columns
+            and value_name in inputs.columns
+        ):
+            raise ValueError(
+                f"DataFrame must have unique_id, ds and {value_name} columns."
+            )
+        if not forecast_context_len:
+            forecast_context_len = self.context_len
+        logging.info("Preprocessing dataframe.")
+        df_sorted = inputs.sort_values(by=["unique_id", "ds"])
+        new_inputs = []
+        uids = []
+        if num_jobs == 1:
+            if verbose:
+                print("Processing dataframe with single process.")
+            for key, group in df_sorted.groupby("unique_id"):
+                inp, uid = process_group(
+                    key,
+                    group,
+                    value_name,
+                    forecast_context_len,
+                )
+                new_inputs.append(inp)
+                uids.append(uid)
+        else:
+            if num_jobs == -1:
+                num_jobs = multiprocessing.cpu_count()
+            if verbose:
+                print("Processing dataframe with multiple processes.")
+            with multiprocessing.Pool(processes=num_jobs) as pool:
+                results = pool.starmap(
+                    process_group,
+                    [
+                        (key, group, value_name, forecast_context_len)
+                        for key, group in df_sorted.groupby("unique_id")
+                    ],
+                )
+            new_inputs, uids = zip(*results)
+        if verbose:
+            print("Finished preprocessing dataframe.")
+        freq_inps = [freq_map(freq)] * len(new_inputs)
+        _, full_forecast = self.forecast(
+            new_inputs, freq=freq_inps, window_size=window_size
+        )
+        if verbose:
+            print("Finished forecasting.")
+        fcst_df = make_future_dataframe(
+            uids=uids,
+            last_times=df_sorted.groupby("unique_id")["ds"].tail(1),
+            h=self.horizon_len,
+            freq=freq,
+        )
+        fcst_df[model_name] = full_forecast[:, 0 : self.horizon_len, 0].reshape(-1, 1)
+
+        if self._model.quantiles is not None:
+            for i, q in enumerate(self._model.quantiles):
+                q_col = f"{model_name}-q-{q}"
+                fcst_df[q_col] = full_forecast[:, 0 : self.horizon_len, 1 + i].reshape(
+                    -1, 1
+                )
+                if q == 0.5:
+                    fcst_df[model_name] = fcst_df[q_col]
+        logging.info("Finished creating output dataframe.")
+        return fcst_df
diff --git a/sktime/libs/timesfm/xreg_lib.py b/sktime/libs/timesfm/xreg_lib.py
new file mode 100644
index 000000000..98b59dc87
--- /dev/null
+++ b/sktime/libs/timesfm/xreg_lib.py
@@ -0,0 +1,361 @@
+# Copyright 2024 The Google Research Authors.
+#
+# 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 sktime.utils.dependencies import _check_soft_dependencies
+
+if _check_soft_dependencies("jax", severity="none"):
+    import jax
+    import jax.numpy as jnp
+
+import numpy as np
+from sklearn import preprocessing
+
+_TOL = 1e-6
+
+
+def _unnest(nested):
+    return np.array(list(itertools.chain.from_iterable(nested)))
+
+
+def _repeat(elements, counts):
+    return np.array(
+        list(itertools.chain.from_iterable(map(itertools.repeat, elements, counts)))
+    )
+
+
+def _to_padded_jax_array(x):
+    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}")
+
+
+class BatchedInContextXRegBase:
+    """BatchedInContextXRegBase."""
+
+    def __init__(
+        self,
+        targets,
+        train_lens,
+        test_lens,
+        train_dynamic_numerical_covariates=None,
+        train_dynamic_categorical_covariates=None,
+        test_dynamic_numerical_covariates=None,
+        test_dynamic_categorical_covariates=None,
+        static_numerical_covariates=None,
+        static_categorical_covariates=None,
+    ):
+        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=False):
+        # 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 "
+                        f"{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="first",
+        use_intercept=True,
+        assert_covariates=False,
+        assert_covariate_shapes=False,
+    ):
+        """create_covariate_matrix."""
+        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):
+        """Fit."""
+        raise NotImplementedError("Fit is not implemented.")
+
+
+class BatchedInContextXRegLinear(BatchedInContextXRegBase):
+    """Linear in-context regression model."""
+
+    def fit(
+        self,
+        ridge=0.0,
+        one_hot_encoder_drop="first",
+        use_intercept=True,
+        force_on_cpu=False,
+        max_rows_per_col=0,
+        max_rows_per_col_sample_seed=42,
+        debug_info=False,
+        assert_covariates=False,
+        assert_covariate_shapes=False,
+    ):
+        """Fit."""
+        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
diff --git a/sktime/utils/_estimator_html_repr.py b/sktime/utils/_estimator_html_repr.py
index fabbc72c0..152e2579d 100644
--- a/sktime/utils/_estimator_html_repr.py
+++ b/sktime/utils/_estimator_html_repr.py
@@ -1,10 +1,11 @@
 """Functionality to represent instance of BaseObject as html."""
-# based on the slearn module of the same name
+# based on the sklearn module of the same name
 
 import html
 import importlib
 import uuid
 from contextlib import closing
+from inspect import isclass
 from io import StringIO
 from pathlib import Path
 from string import Template
@@ -113,15 +114,20 @@ def _get_visual_block(base_object):
     elif base_object is None:
         return _VisualBlock("single", base_object, names="None", name_details="None")
 
-    # check if base_object looks like a meta base_object wraps base_object
-    if hasattr(base_object, "get_params"):
-        base_objects = []
-        for key, value in base_object.get_params().items():
-            # Only look at the BaseObjects in the first layer
-            if "__" not in key and hasattr(value, "get_params"):
-                base_objects.append(value)
-        if len(base_objects):
-            return _VisualBlock("parallel", base_objects, names=None)
+    # check if estimator looks like a meta estimator (wraps estimators)
+    if hasattr(base_object, "get_params") and not isclass(base_object):
+        base_objects = [
+            (key, est)
+            for key, est in base_object.get_params(deep=False).items()
+            if hasattr(est, "get_params") and hasattr(est, "fit") and not isclass(est)
+        ]
+        if base_objects:
+            return _VisualBlock(
+                "parallel",
+                [est for _, est in base_objects],
+                names=[f"{key}: {est.__class__.__name__}" for key, est in base_objects],
+                name_details=[str(est) for _, est in base_objects],
+            )
 
     return _VisualBlock(
         "single",
@@ -236,16 +242,34 @@ def _object_html_repr(base_object):
         return html_output
 
 
+def _get_reduced_path(input_path_string):
+    """Remove submodules starting with an underscore to get a reduced path string."""
+    substrings = input_path_string.split(".")
+
+    index_to_remove = None
+    for i, substring in enumerate(substrings):
+        if substring.startswith("_"):
+            index_to_remove = i
+            break
+
+    if index_to_remove is not None:
+        substrings = substrings[:index_to_remove] + substrings[-1:]
+
+    result_string = ".".join(substrings)
+
+    return result_string
+
+
 class _HTMLDocumentationLinkMixin:
     """Mixin class allowing to generate a link to the API documentation.
 
     This mixin relies on three attributes:
-    - `_doc_link_module`: it corresponds to the root module (e.g. `sklearn`). Using this
-      mixin, the default value is `sklearn`.
+    - `_doc_link_module`: it corresponds to the root module (e.g. `sktime`). Using this
+      mixin, the default value is `sktime`.
     - `_doc_link_template`: it corresponds to the template used to generate the
       link to the API documentation. Using this mixin, the default value is
-      `"https://www.sktime.net/en/v{version_url}/api_reference/auto_generated/
-      {modpath}.html"`.
+      `"https://www.sktime.net/en/v{sktime_version}/api_reference/auto_generated/
+      {path_reduced}.html"`.
     - `_doc_link_url_param_generator`: it corresponds to a function that generates the
       parameters to be used in the template when the estimator module and name are not
       sufficient.
@@ -269,7 +293,7 @@ class _HTMLDocumentationLinkMixin:
             "__doc_link_template",
             (
                 f"https://www.sktime.net/en/v{sktime_version}"
-                "/api_reference/auto_generated/{modpath}.html"
+                "/api_reference/auto_generated/{reduced_path}.html"
             ),
         )
 
@@ -294,10 +318,10 @@ class _HTMLDocumentationLinkMixin:
             return ""
 
         if self._doc_link_url_param_generator is None:
-            modclass = self.__class__
-            modpath = str(modclass)[8:-2]
+            modpath = str(self.__class__)[8:-2]
+            reduced_path = _get_reduced_path(modpath)
 
-            return self._doc_link_template.format(modpath=modpath)
+            return self._doc_link_template.format(reduced_path=reduced_path)
         return self._doc_link_template.format(
             **self._doc_link_url_param_generator(self)
         )
diff --git a/sktime/utils/_maint/_show_versions.py b/sktime/utils/_maint/_show_versions.py
index cc2c621ce..2cddc427d 100644
--- a/sktime/utils/_maint/_show_versions.py
+++ b/sktime/utils/_maint/_show_versions.py
@@ -52,28 +52,53 @@ DEFAULT_DEPS_TO_SHOW = [
     "tslearn",
     "torch",
     "tensorflow",
-    "tensorflow_probability",
 ]
 
 
-def _get_deps_info(deps=None):
+def _get_deps_info(deps=None, source="distributions"):
     """Overview of the installed version of main dependencies.
 
     Parameters
     ----------
-    deps : optional, list of strings with import names
-        if None, behaves as deps = ["sktime"]
+    deps : optional, list of strings with package names
+        if None, behaves as deps = ["sktime"].
+
+    source : str, optional one of "distributions" (default) or "import"
+        source of version information
+
+        * "distributions" - uses importlib.distributions. In this case,
+          strings in deps are assumed to be PEP 440 package strings,
+          e.g., scikit-learn, not sklearn.
+        * "import" - uses the __version__ attribute of the module.
+          In this case, strings in deps are assumed to be import names,
+          e.g., sklearn, not scikit-learn.
 
     Returns
     -------
     deps_info: dict
         version information on libraries in `deps`
-        keys are import names, values are PEP 440 version strings
+        keys are package names, import names if source is "import",
+        and PEP 440 package strings if source is "distributions";
+        values are PEP 440 version strings
         of the import as present in the current python environment
     """
     if deps is None:
         deps = ["sktime"]
 
+    if source == "distributions":
+        from sktime.utils.dependencies._dependencies import _get_installed_packages
+
+        KEY_ALIAS = {"sklearn": "scikit-learn", "skbase": "scikit-base"}
+
+        pkgs = _get_installed_packages()
+
+        deps_info = {}
+        for modname in deps:
+            pkg_name = KEY_ALIAS.get(modname, modname)
+            deps_info[modname] = pkgs.get(pkg_name, None)
+
+        return deps_info
+
     def get_version(module):
         return getattr(module, "__version__", None)
 


Test patch:
diff --git a/sktime/tests/test_softdeps.py b/sktime/tests/test_softdeps.py
index 6f0c7bb7a..599dd7fb4 100644
--- a/sktime/tests/test_softdeps.py
+++ b/sktime/tests/test_softdeps.py
@@ -147,7 +147,7 @@ def _coerce_list_of_str(obj):
 
 
 def _get_soft_deps(est):
-    """Return soft dependencies of an estimator, as list of str and its alias."""
+    """Return soft dependencies of an estimator, as list of str."""
     softdeps = est.get_class_tag("python_dependencies", None)
     softdeps = _coerce_list_of_str(softdeps)
     if softdeps is None:
diff --git a/sktime/utils/tests/test_estimator_html_repr.py b/sktime/utils/tests/test_estimator_html_repr.py
index 320f75af9..8af09c617 100644
--- a/sktime/utils/tests/test_estimator_html_repr.py
+++ b/sktime/utils/tests/test_estimator_html_repr.py
@@ -22,12 +22,12 @@ def test_html_documentation_link_mixin_sktime(mock_version):
         # other branches than `main` (that is "dev").
         assert (
             mixin._doc_link_template == f"https://www.sktime.net/en/v{sktime_version}"
-            "/api_reference/auto_generated/{modpath}.html"
+            "/api_reference/auto_generated/{reduced_path}.html"
         )
         assert (
             mixin._get_doc_link() == f"https://www.sktime.net/en/v{sktime_version}"
             "/api_reference/auto_generated/"
-            "sktime.utils._estimator_html_repr._HTMLDocumentationLinkMixin.html"
+            "sktime.utils._HTMLDocumentationLinkMixin.html"
         )
 
 
@@ -37,10 +37,10 @@ def test_html_documentation_link_mixin_get_doc_link():
 
     # if we set `_doc_link`, then we expect to infer a module and name for the estimator
     mixin._doc_link_module = "sktime"
-    mixin._doc_link_template = "https://website.com/{modpath}.html"
+    mixin._doc_link_template = "https://website.com/{reduced_path}.html"
     assert (
         mixin._get_doc_link() == "https://website.com/"
-        "sktime.utils._estimator_html_repr._HTMLDocumentationLinkMixin.html"
+        "sktime.utils._HTMLDocumentationLinkMixin.html"
     )