Fix

Signed-off-by: Justin Chu <[email protected]>

Author: justinchuby

src/onnx_ir/__init__.py

@@ -81,6 +81,8 @@
     # IO
     "load",
     "save",
+    # Flags
+    "DEBUG",
 ]
 
 from onnx_ir import convenience, external_data, passes, serde, tape, traversal
@@ -147,12 +149,14 @@
 )
 from onnx_ir.serde import TensorProtoTensor, from_onnx_text, from_proto, to_proto
 
+DEBUG = False
 
 def __set_module() -> None:
     """Set the module of all functions in this module to this public module."""
     global_dict = globals()
     for name in __all__:
-        global_dict[name].__module__ = __name__
-
+        if hasattr(global_dict[name], "__module__"):
+            # Set the module of the function to this module
+            global_dict[name].__module__ = __name__
 
 __set_module()

src/onnx_ir/_convenience/__init__.py

@@ -355,7 +355,6 @@ def replace_nodes_and_values(
         old_values: The values to replace.
         new_values: The values to replace with.
     """
-
     for old_value, new_value in zip(old_values, new_values):
         # Propagate relevant info from old value to new value
         # TODO(Rama): Perhaps this should be a separate utility function. Also, consider

src/onnx_ir/_core.py

@@ -1279,7 +1279,7 @@ def _short_tensor_str_for_node(x: Value) -> str:
 
 
 def _normalize_domain(domain: str) -> str:
-    """Normalize 'ai.onnx' to ''"""
+    """Normalize 'ai.onnx' to ''."""
     return "" if domain == "ai.onnx" else domain
 
 
@@ -1709,7 +1709,7 @@ def dtype(self, value: _enums.DataType) -> None:
 
     @property
     def elem_type(self) -> _enums.DataType:
-        """Return the element type of the tensor type"""
+        """Return the element type of the tensor type."""
         return self.dtype
 
     def __hash__(self) -> int:
@@ -2099,7 +2099,6 @@ def Input(
 
     This is equivalent to calling ``Value(name=name, shape=shape, type=type, doc_string=doc_string)``.
     """
-
     # NOTE: The function name is capitalized to maintain API backward compatibility.
 
     return Value(name=name, shape=shape, type=type, doc_string=doc_string)

src/onnx_ir/_graph_containers.py

@@ -14,6 +14,8 @@
 import collections
 from typing import TYPE_CHECKING, Iterable, SupportsIndex
 
+import onnx_ir
+
 if TYPE_CHECKING:
     from onnx_ir import _core
 

src/onnx_ir/external_data.py

@@ -35,8 +35,7 @@
 
 @dataclasses.dataclass
 class _ExternalDataInfo:
-    """
-    A class that stores information about a tensor that is to be stored as external data.
+    """A class that stores information about a tensor that is to be stored as external data.
 
     Attributes:
         name: The name of the tensor that is to be stored as external data.

src/onnx_ir/passes/_pass_infra.py

@@ -71,7 +71,6 @@ class PassResult:
 class PassBase(abc.ABC):
     """Base class for all passes.
 
-
     ``in_place`` and ``changes_input`` properties and what they mean:
 
     +------------+------------------+----------------------------+

src/onnx_ir/passes/common/_c_api_utils.py

@@ -34,7 +34,6 @@ def call_onnx_api(func: Callable[[onnx.ModelProto], _R], model: ir.Model) -> _R:
     Returns:
         The resulting ModelProto that contains the result of the API call.
     """
-
     # Store the original initializer values so they can be restored
     initializer_values = tuple(model.graph.initializers.values())
     tensors = {v.name: v.const_value for v in initializer_values}

src/onnx_ir/passes/common/inliner.py

@@ -1,6 +1,6 @@
 # Copyright (c) ONNX Project Contributors
 # SPDX-License-Identifier: Apache-2.0
-"""Implementation of an inliner for onnx_ir"""
+"""Implementation of an inliner for onnx_ir."""
 
 from __future__ import annotations
 
@@ -11,8 +11,8 @@
 from collections import defaultdict
 from typing import Iterable, List, Sequence, Tuple
 
-import onnx_ir.convenience as _ir_convenience
 import onnx_ir as ir
+import onnx_ir.convenience as _ir_convenience
 
 # A replacement for a node specifies a list of nodes that replaces the original node,
 # and a list of values that replaces the original node's outputs.

src/onnx_ir/serde_test.py

@@ -9,8 +9,7 @@
 import parameterized
 
 import onnx_ir as ir
-from onnx_ir import _version_utils
-from onnx_ir import serde
+from onnx_ir import _version_utils, serde
 
 
 class ConvenienceFunctionsTest(unittest.TestCase):

tests/ir/public_api_test.py

@@ -1,5 +1,3 @@
-# Copyright (c) ONNX Project Contributors
-# SPDX-License-Identifier: Apache-2.0
 # Adapted from
 # https://github.com/pytorch/pytorch/blob/b505e8647547f029d0f7df408ee5f2968f757f89/test/test_public_bindings.py#L523
 # Original code PyTorch license https://github.com/pytorch/pytorch/blob/main/LICENSE
@@ -21,6 +19,7 @@
 
 def _find_all_importables(pkg):
     """Find all importables in the project.
+
     Return them in order.
     """
     return sorted(
@@ -34,6 +33,7 @@ def _find_all_importables(pkg):
 
 def _discover_path_importables(pkg_path: os.PathLike, pkg_name: str) -> Iterable[str]:
     """Yield all importables under a given path and package.
+
     This is like pkgutil.walk_packages, but does *not* skip over namespace
     packages. Taken from https://stackoverflow.com/questions/41203765/init-py-required-for-pkgutil-walk-packages-in-python3
     """
@@ -155,7 +155,8 @@ class TestPublicApiNamespace(unittest.TestCase):
     tested_modules = (IR_NAMESPACE, *(_find_all_importables(onnx_ir)))
 
     def test_correct_module_names(self):
-        """
+        """Test module names are correct.
+
         An API is considered public, if  its  `__module__` starts with `onnx_ir`
         and there is no name in `__module__` or the object itself that starts with "_".
         Each public package should either: