From 157f1e04fbb153353bdd3e600a0e86f9f922bc21 Mon Sep 17 00:00:00 2001
From: Eta <esyra@coreweave.com>
Date: Tue, 5 Aug 2025 11:50:57 -0500
Subject: [PATCH] fix(torch_compat): Replace type annotation not in old `torch`
 versions

---
 CHANGELOG.md               |  8 ++++++++
 tensorizer/_version.py     |  2 +-
 tensorizer/torch_compat.py | 12 +++++++-----
 tests/test_torch_compat.py |  4 +++-
 4 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index c4d2b30..28b29ee 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,13 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [2.11.1] - 2025-08-05
+
+### Fixed
+
+- Fixed an incompatibility between `tensorizer.torch_compat` and `torch<2.7.1`
+  - Thank you to `stuart-mv` for catching this
+
 ## [2.11.0] - 2025-08-04
 
 ### Added
@@ -485,6 +492,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   - `get_gpu_name`
   - `no_init_or_tensor`
 
+[2.11.1]: https://github.com/coreweave/tensorizer/compare/v2.11.0...v2.11.1
 [2.11.0]: https://github.com/coreweave/tensorizer/compare/v2.10.1...v2.11.0
 [2.10.1]: https://github.com/coreweave/tensorizer/compare/v2.10.0...v2.10.1
 [2.10.0]: https://github.com/coreweave/tensorizer/compare/v2.9.3...v2.10.0
diff --git a/tensorizer/_version.py b/tensorizer/_version.py
index 7f6646a..200c236 100644
--- a/tensorizer/_version.py
+++ b/tensorizer/_version.py
@@ -1 +1 @@
-__version__ = "2.11.0"
+__version__ = "2.11.1"
diff --git a/tensorizer/torch_compat.py b/tensorizer/torch_compat.py
index 8345dc4..f0a075a 100644
--- a/tensorizer/torch_compat.py
+++ b/tensorizer/torch_compat.py
@@ -65,9 +65,11 @@
     int,
 ]
 
+_FileLike: "typing.TypeAlias" = Union[str, os.PathLike[str], typing.IO[bytes]]
+
 _wrapper_file_obj_type: "typing.TypeAlias" = Union[
     _tensorizer_file_obj_type,
-    Callable[[torch.types.FileLike], _tensorizer_file_obj_type],
+    Callable[[_FileLike], _tensorizer_file_obj_type],
 ]
 
 _save_func_type: "typing.TypeAlias" = Callable[
@@ -397,7 +399,7 @@ def _pickle_attr(name):
 _ORIG_TORCH_LOAD: Final[callable] = torch.load
 
 
-def _infer_tensor_ext_name(f: torch.types.FileLike):
+def _infer_tensor_ext_name(f: _FileLike):
     if isinstance(f, io.BytesIO):
         logger.warning(
             "Cannot infer .tensors location from io.BytesIO;"
@@ -418,7 +420,7 @@ def _infer_tensor_ext_name(f: torch.types.FileLike):
 
 @contextlib.contextmanager
 def _contextual_torch_filename(
-    f: torch.types.FileLike,
+    f: _FileLike,
     filename_ctx_var: ContextVar[Optional[_wrapper_file_obj_type]],
 ):
     if filename_ctx_var.get() is None:
@@ -462,7 +464,7 @@ def _contextual_torch_filename(
 @functools.wraps(_ORIG_TORCH_SAVE)
 def _save_wrapper(
     obj: object,
-    f: torch.types.FileLike,
+    f: _FileLike,
     pickle_module: Any = pickle,
     *args,
     **kwargs,
@@ -489,7 +491,7 @@ def _save_wrapper(
 
 @functools.wraps(_ORIG_TORCH_LOAD)
 def _load_wrapper(
-    f: torch.types.FileLike,
+    f: _FileLike,
     map_location: torch.serialization.MAP_LOCATION = None,
     pickle_module: Any = _LOAD_WRAPPER_DEFAULT_MODULE,
     *args,
diff --git a/tests/test_torch_compat.py b/tests/test_torch_compat.py
index 62d343b..6feec72 100644
--- a/tests/test_torch_compat.py
+++ b/tests/test_torch_compat.py
@@ -402,7 +402,9 @@ def test_name_callback(self):
             ".pt.tensors.dynamic"
         )
 
-        def path_callback(f: torch.types.FileLike) -> io.BufferedIOBase:
+        def path_callback(
+            f: tensorizer.torch_compat._FileLike,
+        ) -> io.BufferedIOBase:
             # Test with an exotic function that returns a pre-opened
             # stream dynamically, based on the input file's name
             _path = Path(f).with_suffix(".pt.tensors.dynamic")