加入 infinicore.use_ntops

Author: voltjia

python/infinicore/__init__.py

@@ -24,6 +24,7 @@
     short,
     uint8,
 )
+from infinicore.ntops import use_ntops
 from infinicore.ops.matmul import matmul
 from infinicore.ops.rearrange import rearrange
 from infinicore.tensor import (
@@ -62,6 +63,8 @@
     "long",
     "short",
     "uint8",
+    # `ntops` integration.
+    "use_ntops",
     # Operations.
     "matmul",
     "rearrange",

python/infinicore/ntops.py

@@ -0,0 +1,55 @@
+import sys
+
+import infinicore
+
+
+def use_ntops():
+    import ntops
+
+    return _TemporaryAttributes(
+        (("ntops.torch.torch", infinicore),)
+        + tuple(
+            (f"infinicore.{op_name}", getattr(ntops.torch, op_name))
+            for op_name in ntops.torch.__all__
+        )
+    )
+
+
+class _TemporaryAttributes:
+    def __init__(self, attribute_mappings):
+        self._attribute_mappings = attribute_mappings
+
+        self._original_values = {}
+
+    def __enter__(self):
+        for attr_path, new_value in self._attribute_mappings:
+            parent, attr_name = self._resolve_path(attr_path)
+
+            try:
+                self._original_values[attr_path] = getattr(parent, attr_name)
+            except AttributeError:
+                pass
+
+            setattr(parent, attr_name, new_value)
+
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        for attr_path, _ in self._attribute_mappings:
+            parent, attr_name = self._resolve_path(attr_path)
+
+            if attr_path in self._original_values:
+                setattr(parent, attr_name, self._original_values[attr_path])
+            else:
+                delattr(parent, attr_name)
+
+    @staticmethod
+    def _resolve_path(path):
+        *parent_parts, attr_name = path.split(".")
+
+        curr = sys.modules[parent_parts[0]]
+
+        for part in parent_parts[1:]:
+            curr = getattr(curr, part)
+
+        return curr, attr_name