From 47be457206fc50d95c6e59ccffc5a34ceb450291 Mon Sep 17 00:00:00 2001
From: Oleksandr Shchur <shchuro@amazon.com>
Date: Fri, 20 Feb 2026 17:06:04 +0000
Subject: [PATCH] Use typing_extensions for NotRequired

---
 pyproject.toml                  | 1 +
 src/chronos/chronos2/dataset.py | 7 ++++++-
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index d9e7117..d21994f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -20,6 +20,7 @@ dependencies = [
   "numpy>=1.21,<3",
   "einops>=0.7.0,<1",
   "scikit-learn>=1.6.0,<2",
+  "typing_extensions~=4.0",
 ]
 classifiers = [
   "Programming Language :: Python :: 3",
diff --git a/src/chronos/chronos2/dataset.py b/src/chronos/chronos2/dataset.py
index 20d2528..dbde6fe 100644
--- a/src/chronos/chronos2/dataset.py
+++ b/src/chronos/chronos2/dataset.py
@@ -5,7 +5,7 @@
 
 import math
 from enum import Enum
-from typing import TYPE_CHECKING, Any, Iterable, Iterator, Mapping, NotRequired, Sequence, TypeAlias, TypedDict, cast
+from typing import TYPE_CHECKING, Any, Iterable, Iterator, Mapping, Sequence, TypeAlias, TypedDict, cast
 
 import numpy as np
 import torch
@@ -16,6 +16,11 @@ if TYPE_CHECKING:
     import datasets
     import fev
 
+try:
+    from typing import NotRequired  # Python 3.11+
+except ImportError:
+    from typing_extensions import NotRequired
+
 
 TensorOrArray: TypeAlias = torch.Tensor | np.ndarray