diff --git a/README.md b/README.md
index c29d978..fbe0d9f 100644
--- a/README.md
+++ b/README.md
@@ -78,18 +78,23 @@ poetry install -E  torch
 
 After than you can run the timesfm under `poetry shell` or do `poetry run python3 ...`.
 
-**Note**: 
+**Additional Note**: 
 
-1. Running the provided benchmarks would require additional dependencies.
-Please see the `experiments` section fro more instructions.
+If you plan to use the **`forecast_with_covariates`** function (which requires external regressors), 
+you need to install **JAX** and **jaxlib**. If you installed the base version of TimesFM (`torch`), you must manually install the dependencies for **`forecast_with_covariates`** support:
+```
+pip install jax jaxlib
+```
 
-2. The dependency `lingvo` does not support ARM architectures, and the code is not working for machines with Apple silicon. We are aware of this issue and are working on a solution. Stay tuned.
+**Why is this needed?**  
+The `forecast_with_covariates` method relies on the `xreg_lib` module, which depends on JAX and jaxlib. If these packages are not installed, 
+calling `forecast_with_covariates` will raise an error. However, due to a lazy import mechanism, `xreg_lib` (and hence JAX/jaxlib) is not needed for standard `forecast` calls.
 
 ### Notes
 
 1. Running the provided benchmarks would require additional dependencies. Please see the `experiments` folder.
 
-2. The dependency `lingvo` does not support ARM architectures, and the PAX version is not working for machines with Apple silicon.
+2. The dependency `lingvo` does not support ARM architectures, and the code is not working for machines with Apple silicon. We are aware of this issue and are working on a solution. Stay tuned.
 
 ### Install from PyPI (and publish)
 
@@ -241,6 +246,12 @@ forecast_df = tfm.forecast_on_df(
 
 We now have an external regressors library on top of TimesFM that can support static covariates as well as dynamic covariates available in the future. We have an usage example in [notebooks/covariates.ipynb](https://github.com/google-research/timesfm/blob/master/notebooks/covariates.ipynb).
 
+If you plan to use the **`forecast_with_covariates`** on timesfm `torch` version, you need to install **JAX** and **jaxlib**. 
+You must manually install the dependencies for **`forecast_with_covariates`** support:
+```
+pip install jax jaxlib
+```
+
 Let's take a toy example of forecasting sales for a grocery store: 
 
 **Task:** Given the observed the daily sales of this week (7 days), forecast the daily sales of next week (7 days).
diff --git a/pyproject.toml b/pyproject.toml
index d2084f0..09a0a18 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -42,7 +42,7 @@ absl-py = ">=1.4.0"
 
 [tool.poetry.extras]
 pax = ["paxml", "lingvo", "jax", "jaxlib"]
-torch = ["torch",  "jax", "jaxlib"]  # jax & jaxlib are already in pax
+torch = ["torch"] 
 
 [tool.poetry.dependencies.paxml]
 version = ">=1.4.0"
diff --git a/src/timesfm/timesfm_base.py b/src/timesfm/timesfm_base.py
index 088cfab..b0a622a 100644
--- a/src/timesfm/timesfm_base.py
+++ b/src/timesfm/timesfm_base.py
@@ -17,17 +17,20 @@ import collections
 import dataclasses
 import logging
 import multiprocessing
-from typing import Any, Literal, Sequence
+from typing import Any, Literal, Sequence, TYPE_CHECKING
 
 import numpy as np
 import pandas as pd
 
 from utilsforecast.processing import make_future_dataframe
 
-from . import xreg_lib
-
-Category = xreg_lib.Category
-XRegMode = xreg_lib.XRegMode
+if TYPE_CHECKING:
+    from . import xreg_lib
+    Category = xreg_lib.Category
+    XRegMode = xreg_lib.XRegMode
+else:
+    Category = int | str
+    XRegMode = str
 
 _TOL = 1e-6
 DEFAULT_QUANTILES = (0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9)
@@ -42,7 +45,7 @@ def moving_average(arr, window_size):
   """Calculates the moving average using NumPy's convolution function."""
   # 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") /
+  smoothed_arr = (np.convolve(arr_padded, np.ones(window_size), "valid") / 
                   window_size)
   return [smoothed_arr, arr - smoothed_arr]
 
@@ -464,6 +467,8 @@ class TimesFmBase:
       the outputs of the xreg.
     """
 
+    from . import xreg_lib
+
     # Verify and bookkeep covariates.
     if not (dynamic_numerical_covariates or dynamic_categorical_covariates or
             static_numerical_covariates or static_categorical_covariates):