Update base for Update on "[NJT] Allow construction of NJT within graph using offsets from inputs"

soulitzer · soulitzer · commit cd52984c17c6 · 2024-04-22T14:57:40.000-04:00
Creating symbolic nested ints within the graph is difficult. Using unbacked symints should solve the most important(?) cases in the mean time. See #118446 Known gaps: - creating NJT from intermediate offsets (offsets created within the graph, as opposed to being offsets passed in as inputs) - when the same offsets is also passed in as a input to the graph. We are not smart enough to realize that the offsets from that input is the same and therefore would fail when the sizes are compare ("s0 cannot be compared with u0") cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang [ghstack-poisoned]
diff --git a/torch/onnx/_internal/fx/patcher.py b/torch/onnx/_internal/fx/patcher.py
@@ -1,17 +1,18 @@
 import copy
+import functools
 import io
 from typing import List, Union
-import functools
 
 import torch
 
+
 # TODO: Remove after https://github.com/huggingface/safetensors/pull/318
-@functools.cache
+@functools.lru_cache(None)
 def has_safetensors_and_transformers():
     try:
         # safetensors is not an exporter requirement, but needed for some huggingface models
         import safetensors  # type: ignore[import]  # noqa: F401
-        import transformers  # type: ignore[import]
+        import transformers  # type: ignore[import]  # noqa: F401
 
         from safetensors import torch as safetensors_torch  # noqa: F401
 
@@ -66,6 +67,8 @@ def torch_load_wrapper(f, *args, **kwargs):
         self.torch_load_wrapper = torch_load_wrapper
 
         if has_safetensors_and_transformers():
+            import safetensors
+            import transformers
 
             def safetensors_load_file_wrapper(filename, device="cpu"):
                 # Record path for later serialization into ONNX proto
@@ -114,6 +117,9 @@ def __enter__(self):
         torch.fx._symbolic_trace._wrapped_methods_to_patch = desired_wrapped_methods
 
         if has_safetensors_and_transformers():
+            import safetensors
+            import transformers
+
             safetensors.torch.load_file = self.safetensors_torch_load_file_wrapper
             transformers.modeling_utils.safe_load_file = (
                 self.safetensors_torch_load_file_wrapper
@@ -125,6 +131,9 @@ def __exit__(self, exc_type, exc_value, traceback):
             self.torch_fx__symbolic_trace__wrapped_methods_to_patch
         )
         if has_safetensors_and_transformers():
+            import safetensors
+            import transformers
+
             safetensors.torch.load_file = self.safetensors_torch_load_file
             transformers.modeling_utils.safe_load_file = (
                 self.transformers_modeling_utils_safe_load_file
