fix(torch_compat): Replace type annotation not in old torch versions

Author: Eta0

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

tensorizer/_version.py

@@ -1 +1 @@
-__version__ = "2.11.0"
+__version__ = "2.11.1"

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,

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")