tico 0.1.0.dev250714__py3-none-any.whl → 0.1.0.dev250715__py3-none-any.whl

tico/__init__.py

@@ -21,7 +21,7 @@ from tico.config import CompileConfigV1, get_default_config
 from tico.utils.convert import convert, convert_from_exported_program, convert_from_pt2
 
 # THIS LINE IS AUTOMATICALLY GENERATED BY setup.py
-__version__ = "0.1.0.dev250714"
+__version__ = "0.1.0.dev250715"
 
 MINIMUM_SUPPORTED_VERSION = "2.5.0"
 SECURE_TORCH_VERSION = "2.6.0"

tico/utils/pytree_utils.py

@@ -0,0 +1,134 @@
+import threading
+
+import torch
+from packaging.version import Version
+
+from tico.utils import logging
+from tico.utils.installed_packages import is_transformers_installed
+
+__all__ = ["register_dynamic_cache"]
+
+
+def register_dynamic_cache():
+    PyTreeRegistryHelper().register_dynamic_cache()
+
+
+class PyTreeRegistryHelper:
+    """
+    Thread-safe singleton helper class for registering custom PyTree nodes.
+
+    This class provides functionality to register DynamicCache as a PyTree node
+    for torch.export compatibility. This registration is only needed for
+    transformers versions below 4.50.0.
+
+    Thread Safety:
+    - Uses a class-level threading.Lock() to ensure thread-safe singleton instantiation
+    - Uses the same lock to protect the registration process from concurrent calls
+    """
+
+    _instance = None  # Class variable to hold the singleton instance
+    _has_called = False  # Flag to track if registration has been performed
+    _lock = threading.Lock()  # Class-level lock for thread-safe operations
+
+    def __init__(self):
+        """Private constructor to prevent direct instantiation"""
+        pass
+
+    def __new__(cls, *args, **kwargs):
+        """
+        Thread-safe singleton instance creation using double-checked locking pattern.
+
+        Returns:
+            PyTreeRegistryHelper: The singleton instance of this class
+        """
+        if not cls._instance:
+            with cls._lock:  # Acquire lock for thread-safe instantiation
+                if not cls._instance:  # Double-check after acquiring lock
+                    cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def register_dynamic_cache(self):
+        """
+        Registers DynamicCache as a PyTree node for torch.export compatibility.
+
+        This method is thread-safe and idempotent - it will only perform the
+        registration once, even if called multiple times from different threads.
+
+        Note:
+            This registration is only needed for transformers versions below 4.50.0.
+
+        Raises:
+            ImportError: If transformers package is not installed
+        """
+        with self._lock:  # Acquire lock for thread-safe registration
+            if self.__class__._has_called:
+                logger = logging.getLogger(__name__)
+                logger.debug("register_dynamic_cache already called, skipping")
+                return
+
+            self.__class__._has_called = True
+            logger = logging.getLogger(__name__)
+            logger.info("Registering DynamicCache PyTree node")
+
+        if not is_transformers_installed:  # type: ignore[truthy-function]
+            raise ImportError("transformers package is not installed")
+
+        import transformers
+
+        HAS_TRANSFORMERS_LESS_4_50_0 = Version(transformers.__version__) < Version(
+            "4.50.0"
+        )
+        if not HAS_TRANSFORMERS_LESS_4_50_0:
+            return
+
+        from transformers.cache_utils import DynamicCache
+
+        def _flatten_dynamic_cache(dynamic_cache: DynamicCache):
+            if not isinstance(dynamic_cache, DynamicCache):
+                raise RuntimeError(
+                    "This pytree flattening function should only be applied to DynamicCache"
+                )
+            HAS_TORCH_2_6_0 = Version(torch.__version__) >= Version("2.6.0")
+            if not HAS_TORCH_2_6_0:
+                logger = logging.getLogger(__name__)
+                logger.warning_once(
+                    "DynamicCache + torch.export is tested on torch 2.6.0+ and may not work on earlier versions."
+                )
+            dictionary = {
+                "key_cache": getattr(dynamic_cache, "key_cache"),
+                "value_cache": getattr(dynamic_cache, "value_cache"),
+            }
+            return torch.utils._pytree._dict_flatten(dictionary)
+
+        def _flatten_with_keys_dynamic_cache(dynamic_cache: DynamicCache):
+            dictionary = {
+                "key_cache": getattr(dynamic_cache, "key_cache"),
+                "value_cache": getattr(dynamic_cache, "value_cache"),
+            }
+            return torch.utils._pytree._dict_flatten_with_keys(dictionary)
+
+        def _unflatten_dynamic_cache(values, context: torch.utils._pytree.Context):
+            dictionary = torch.utils._pytree._dict_unflatten(values, context)
+            cache = DynamicCache()
+            for k, v in dictionary.items():
+                setattr(cache, k, v)
+            return cache
+
+        def _flatten_dynamic_cache_for_fx(cache, spec):
+            dictionary = {
+                "key_cache": getattr(cache, "key_cache"),
+                "value_cache": getattr(cache, "value_cache"),
+            }
+            return torch.fx._pytree._dict_flatten_spec(dictionary, spec)
+
+        torch.utils._pytree.register_pytree_node(
+            DynamicCache,
+            _flatten_dynamic_cache,
+            _unflatten_dynamic_cache,
+            serialized_type_name=f"{DynamicCache.__module__}.{DynamicCache.__name__}",
+            flatten_with_keys_fn=_flatten_with_keys_dynamic_cache,
+        )
+        # TODO: This won't be needed in torch 2.7+.
+        torch.fx._pytree.register_pytree_flatten_spec(
+            DynamicCache, _flatten_dynamic_cache_for_fx
+        )

{tico-0.1.0.dev250714.dist-info → tico-0.1.0.dev250715.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tico
-Version: 0.1.0.dev250714
+Version: 0.1.0.dev250715
 Summary: Convert exported Torch module to circle
 Home-page: UNKNOWN
 License: UNKNOWN

{tico-0.1.0.dev250714.dist-info → tico-0.1.0.dev250715.dist-info}/RECORD

@@ -1,4 +1,4 @@
-tico/__init__.py,sha256=xbMeC4-9n5Z8Lu1B6yqbmE10G6wmHdhpue7kD7skYtM,1743
+tico/__init__.py,sha256=MYWB0f9ftIZXXj1q1Sdv4Qn0EGgO27twfOAD_gDGNVQ,1743
 tico/pt2_to_circle.py,sha256=gu3MD4Iqc0zMZcCZ2IT8oGbyj21CTSbT3Rgd9s2B_9A,2767
 tico/config/__init__.py,sha256=xZzCXjZ84qE-CsBi-dfaL05bqpQ3stKKfTXhnrJRyVs,142
 tico/config/base.py,sha256=anwOiJFkUxUi7Cef573JgQcjk6S-FSi6O_TLjYASW-g,1244
@@ -192,6 +192,7 @@ tico/utils/logging.py,sha256=IlbBWscsaHidI0dNqro1HEXAbIcbkR3BD5ukLy2m95k,1286
 tico/utils/model.py,sha256=Uqc92AnJXQ2pbvctS2z2F3Ku3yNrwXZ9O33hZVis7is,1250
 tico/utils/padding.py,sha256=jyNhGmlLZfruWZ6n5hll8RZOFg85iCZP8OJqnHGS97g,3293
 tico/utils/passes.py,sha256=kGmDe__5cPaO6i5EDAoXSVe6yXEoX9hAny4ROb3ZEmQ,2409
+tico/utils/pytree_utils.py,sha256=jrk3N6X6LiUnBCX_gM1K9nywbVAJBVnszlTAgeIeDUc,5219
 tico/utils/register_custom_op.py,sha256=3-Yl6iYmx1qQA2igNHt4hYhQhQMkdPb7gF50LIY8yvc,27350
 tico/utils/serialize.py,sha256=AQXMBOLu-Kg2Rn-qbqsAtHndjZAZIavlKA0QFgJREHM,1420
 tico/utils/trace_decorators.py,sha256=ddLIiKQfSaQrxgF1kNpwjFTQnXENzeSfcr1kuAW4jGI,3221
@@ -201,9 +202,9 @@ tico/utils/mx/__init__.py,sha256=IO6FP_xYbGy0dW0HL26GXD3ouxARaxCK7bz9dn4blPQ,26
 tico/utils/mx/elemwise_ops.py,sha256=V6glyAHsVR1joqpsgnNytatCD_ew92xNWZ19UFDoMTA,10281
 tico/utils/mx/formats.py,sha256=uzNWyu-1onUlwQfX5cZ6fZSUfHMRqorper7_T1k3jfk,3404
 tico/utils/mx/mx_ops.py,sha256=RcfUTYVi-wilGB2sC35OeARdwDqnixv7dG5iyZ-fQT8,8555
-tico-0.1.0.dev250714.dist-info/LICENSE,sha256=kp4JLII7bzRhPb0CPD5XTDZMh22BQ7h3k3B7t8TiSbw,12644
-tico-0.1.0.dev250714.dist-info/METADATA,sha256=5RPA0JUl-L5MOz1_Ix5-8SZSpn628DHjCpNZHPvrO68,8430
-tico-0.1.0.dev250714.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
-tico-0.1.0.dev250714.dist-info/entry_points.txt,sha256=kBKYSS_IYrSXmUYevmmepqIVPScq5vF8ulQRu3I_Zf0,59
-tico-0.1.0.dev250714.dist-info/top_level.txt,sha256=oqs7UPoNSKZEwqsX8B-KAWdQwfAa7i60pbxW_Jk7P3w,5
-tico-0.1.0.dev250714.dist-info/RECORD,,
+tico-0.1.0.dev250715.dist-info/LICENSE,sha256=kp4JLII7bzRhPb0CPD5XTDZMh22BQ7h3k3B7t8TiSbw,12644
+tico-0.1.0.dev250715.dist-info/METADATA,sha256=XNqsTtUt8jSqU2EsY3sm3RJDKwRGerDxGkH9eMXwOQk,8430
+tico-0.1.0.dev250715.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
+tico-0.1.0.dev250715.dist-info/entry_points.txt,sha256=kBKYSS_IYrSXmUYevmmepqIVPScq5vF8ulQRu3I_Zf0,59
+tico-0.1.0.dev250715.dist-info/top_level.txt,sha256=oqs7UPoNSKZEwqsX8B-KAWdQwfAa7i60pbxW_Jk7P3w,5
+tico-0.1.0.dev250715.dist-info/RECORD,,