hyperonnx 1.0.0__tar.gz → 1.0.2__tar.gz

{hyperonnx-1.0.0 → hyperonnx-1.0.2}/.github/workflows/lint.yml

@@ -11,6 +11,7 @@ jobs:
     strategy:
       matrix:
         python-version: ["3.12"]
+        transformers-version: ["4.53.0", "5.1.0"]
     steps:
       - uses: actions/checkout@v5
         with:
@@ -24,6 +25,7 @@ jobs:
           node-version: lts/Jod
       - name: pyright type check
         run: |
-          uv sync --extra=test -U
+          uv sync --dev --extra=cpu
+          uv pip install 'transformers==${{ matrix.transformers-version }}'
           uv run pyright hyperonnx
           uv run pre-commit run -s HEAD^ -o HEAD

{hyperonnx-1.0.0 → hyperonnx-1.0.2}/.github/workflows/unittest.yml

@@ -12,6 +12,7 @@ jobs:
     strategy:
       matrix:
         python-version: ["3.12"]
+        transformers-version: ["4.53.0", "5.1.0"]
     steps:
       - uses: actions/checkout@v5
       - name: Install uv and set the Python version
@@ -20,7 +21,8 @@ jobs:
           python-version: ${{ matrix.python-version }}
       - name: Install dependencies
         run: |
-          uv sync --extra=test -U
+          uv sync --dev --extra=cpu -U
+          uv pip install 'transformers==${{ matrix.transformers-version }}'
       - name: Run pytest and coverage
         run: |
-          uv run pytest --cov=hyperonnx tests
+          uv run --dev --extra=cpu pytest --cov=hyperonnx tests

{hyperonnx-1.0.0 → hyperonnx-1.0.2}/.pre-commit-config.yaml

@@ -9,32 +9,26 @@ repos:
     -   id: check-toml
     -   id: check-yaml
     -   id: check-added-large-files
--   repo: https://github.com/PyCQA/isort
-    rev: 6.0.1
+-   repo: https://github.com/astral-sh/ruff-pre-commit
+    # Ruff version.
+    rev: v0.15.8
     hooks:
-    -   id: isort
-        args: ['--profile=black']
--   repo: https://github.com/psf/black
-    rev: 25.1.0
-    hooks:
-    -   id: black
+    # Run the linter.
+    -   id: ruff-check
+        args: [--fix]
+    # Run the formatter.
+    -   id: ruff-format
 -   repo: https://github.com/codespell-project/codespell
-    rev: v2.4.1
+    rev: v2.4.2
     hooks:
     -   id: codespell
         args: [-w, -I=.github/words_bag.txt]
         exclude: (\.(txt|diff|patch)$)
--   repo: https://github.com/pycqa/flake8
-    rev: 7.3.0
-    hooks:
-    - id: flake8
-      additional_dependencies: [Flake8-pyproject]
-      args:
-      - '--ignore=W503,E203,E231,E241'
 -   repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.17.1
+    rev: v1.19.1
     hooks:
     -   id: mypy
         args:
         - '--ignore-missing-imports'
         - '--exclude=tests.*\.py'
+        exclude: (^tests/)

{hyperonnx-1.0.0 → hyperonnx-1.0.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: hyperonnx
-Version: 1.0.0
+Version: 1.0.2
 Summary: Copyright (C) 2025 The HYPERONNX Authors.
 Author-email: Wenyi Tang <wenyitang@outlook.com>
 Requires-Python: >=3.11
@@ -9,26 +9,29 @@ Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Topic :: Software Development
 License-File: LICENSE
-Requires-Dist: onnxifier>=2.0.0
-Requires-Dist: onnxscript
+Requires-Dist: onnxifier>=2.0.1
+Requires-Dist: onnxscript<0.6.0
 Requires-Dist: onnx>=1.18.0
-Requires-Dist: torch
-Requires-Dist: mypy ; extra == "test"
-Requires-Dist: onnxruntime ; extra == "test"
-Requires-Dist: pre-commit ; extra == "test"
-Requires-Dist: pyright>=1.1.408 ; extra == "test"
-Requires-Dist: pytest ; extra == "test"
-Requires-Dist: pytest-cov ; extra == "test"
-Requires-Dist: torchvision ; extra == "test"
+Requires-Dist: torch<2.11 ; extra == "cpu"
+Requires-Dist: torchvision ; extra == "cpu"
+Requires-Dist: triton-windows ; extra == "cpu" and ( sys_platform == 'win32')
+Requires-Dist: torch<2.11 ; extra == "cuda"
+Requires-Dist: torchvision ; extra == "cuda"
+Requires-Dist: triton-windows ; extra == "cuda" and ( sys_platform == 'win32')
+Requires-Dist: torch<2.11 ; extra == "xpu"
+Requires-Dist: torchvision ; extra == "xpu"
+Requires-Dist: pytorch-triton-xpu ; extra == "xpu"
 Project-URL: Docs, https://loseall.github.io/hyperonnx
 Project-URL: Home, https://github.com/LoSealL/hyperonnx
-Provides-Extra: test
+Provides-Extra: cpu
+Provides-Extra: cuda
+Provides-Extra: xpu
 
 # 🚀 HYPER-ONNX
 
 [中文](./README_CN.md)|[EN](./README.md)
 
-Hyper-ONNX can export pytorch models (`nn.Module`) in a hierachical manner. It can keep the module hier information and make a nested onnx graph. ✨
+Hyper-ONNX can export pytorch models (`nn.Module`) in a hierarchical manner. It can keep the module hier information and make a nested onnx graph. ✨
 
 
 ## 📦 Install
@@ -76,7 +79,7 @@ export_hyper_onnx(
 
 ```python
 from hyperonnx import auto_trace_method
-from hyperonnx.patch import patch_transformers
+from hyperonnx.transformers import patch_transformers
 from transformers import (
     GenerationConfig,
     Qwen2_5OmniThinkerForConditionalGeneration,

{hyperonnx-1.0.0 → hyperonnx-1.0.2}/README.md

@@ -2,7 +2,7 @@
 
 [中文](./README_CN.md)|[EN](./README.md)
 
-Hyper-ONNX can export pytorch models (`nn.Module`) in a hierachical manner. It can keep the module hier information and make a nested onnx graph. ✨
+Hyper-ONNX can export pytorch models (`nn.Module`) in a hierarchical manner. It can keep the module hier information and make a nested onnx graph. ✨
 
 
 ## 📦 Install
@@ -50,7 +50,7 @@ export_hyper_onnx(
 
 ```python
 from hyperonnx import auto_trace_method
-from hyperonnx.patch import patch_transformers
+from hyperonnx.transformers import patch_transformers
 from transformers import (
     GenerationConfig,
     Qwen2_5OmniThinkerForConditionalGeneration,

{hyperonnx-1.0.0 → hyperonnx-1.0.2}/README_CN.md

@@ -50,7 +50,7 @@ export_hyper_onnx(
 
 ```python
 from hyperonnx import auto_trace_method
-from hyperonnx.patch import patch_transformers
+from hyperonnx.transformers import patch_transformers
 from transformers import (
     GenerationConfig,
     Qwen2_5OmniThinkerForConditionalGeneration,

{hyperonnx-1.0.0 → hyperonnx-1.0.2}/hyperonnx/__init__.py

@@ -19,4 +19,4 @@ from .hyper_export import export_hyper_onnx
 from .utils import HYPER_DOMAIN
 
 __all__ = ["auto_trace_method", "export_hyper_onnx", "HYPER_DOMAIN"]
-__version__ = "1.0.0"
+__version__ = "1.0.2"

{hyperonnx-1.0.0 → hyperonnx-1.0.2}/hyperonnx/auto.py

@@ -15,13 +15,13 @@ limitations under the License.
 """
 
 import inspect
+from collections.abc import Collection
 from contextlib import contextmanager
 from functools import wraps
 from io import BytesIO
 from os import PathLike
 from pathlib import Path
 from types import MethodType
-from typing import Collection, Dict, List, Optional, Tuple
 
 from torch import inference_mode
 from torch.nn import Module
@@ -60,8 +60,8 @@ class AutoTraceMethod(Module):
 
     def __init__(self, clsmethod: MethodType, expected_stages: int = 1):
         super().__init__()
-        self.pos_args: List[Tuple[AnyTensor, ...]] = []
-        self.kwargs: List[Dict[str, AnyTensor]] = []
+        self.pos_args: list[tuple[AnyTensor, ...]] = []
+        self.kwargs: list[dict[str, AnyTensor]] = []
         model = clsmethod.__self__
         if not isinstance(model, Module):
             for k, v in model.__dict__.items():
@@ -88,15 +88,15 @@ class AutoTraceMethod(Module):
         self,
         f: str | PathLike | BytesIO,
         *,
-        input_names: Optional[List[str]] = None,
-        output_names: Optional[List[str]] = None,
+        input_names: list[str] | None = None,
+        output_names: list[str] | None = None,
         opset_version: int = 19,
         dynamo: bool = False,
         external_data: bool = False,
-        hiera: Optional[Collection[type[Module]]] = None,
-        module_spec: Optional[Dict[Module, ModuleSpec]] = None,
+        hiera: Collection[type[Module]] | None = None,
+        module_spec: dict[Module, ModuleSpec] | None = None,
         do_optimization: bool = True,
-        external_directory: Optional[str | PathLike] = None,
+        external_directory: str | PathLike | None = None,
     ):
         """Export onnx model according to the traced data."""
         if not self.pos_args and not self.kwargs:

{hyperonnx-1.0.0 → hyperonnx-1.0.2}/hyperonnx/exporter/__init__.py

@@ -15,7 +15,6 @@ limitations under the License.
 """
 
 from contextlib import contextmanager
-from typing import Dict
 
 from torch.nn import Module
 
@@ -26,7 +25,7 @@ from .torchscript import replace_duck_forward
 
 @contextmanager
 def replace_with_duck_module(
-    model: Module, dynamo: bool, module_spec: Dict[Module, ModuleSpec]
+    model: Module, dynamo: bool, module_spec: dict[Module, ModuleSpec]
 ):
     """Replace the forward function of modules in `module_spec` with a duck type.
 

{hyperonnx-1.0.0 → hyperonnx-1.0.2}/hyperonnx/exporter/dynamo.py

@@ -17,8 +17,8 @@ limitations under the License.
 import inspect
 from ast import AST
 from collections import OrderedDict
+from collections.abc import Callable
 from contextlib import contextmanager
-from typing import Callable, Dict
 
 import onnxscript
 import onnxscript.ir._schemas as schemas
@@ -257,7 +257,7 @@ def make_custom_op(module: Module, spec: ModuleSpec):
 
 
 @contextmanager
-def replace_with_custom_op(model: Module, module_spec: Dict[Module, ModuleSpec]):
+def replace_with_custom_op(model: Module, module_spec: dict[Module, ModuleSpec]):
     """Replace the forward function of modules in `module_spec` with a duck type.
 
     It's used to laterly replace the duck type with the embedded onnx functions.

{hyperonnx-1.0.0 → hyperonnx-1.0.2}/hyperonnx/exporter/torchscript.py

@@ -16,7 +16,6 @@ limitations under the License.
 
 from contextlib import contextmanager
 from functools import update_wrapper
-from typing import Dict, List
 
 import torch
 from torch import Tensor
@@ -91,7 +90,7 @@ def make_duck_forward(module_spec: ModuleSpec):
 
 
 @contextmanager
-def replace_duck_forward(model: Module, module_spec: Dict[Module, ModuleSpec]):
+def replace_duck_forward(model: Module, module_spec: dict[Module, ModuleSpec]):
     """Replace the forward function of modules in `module_spec` with a duck type.
 
     It's used to laterly replace the duck type with the embedded onnx functions.
@@ -111,7 +110,7 @@ def replace_duck_forward(model: Module, module_spec: Dict[Module, ModuleSpec]):
         output = spec.get("output")
         return output
 
-    handles: List[RemovableHandle] = []
+    handles: list[RemovableHandle] = []
     try:
         for child in filter(lambda c: c in module_spec, model.modules()):
             if module_spec[child]["status"] == ExportStatus.EXPORTED:

{hyperonnx-1.0.0 → hyperonnx-1.0.2}/hyperonnx/exporter/utils.py

@@ -14,8 +14,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+from collections.abc import Sequence
 from itertools import chain
-from typing import Sequence, Tuple
 
 import torch
 from onnxifier.logger import warning
@@ -24,7 +24,7 @@ from torch import Tensor
 from ..typing import AnyTensor, ModuleSpec
 
 
-def plain_tensor_container(obj: AnyTensor) -> Tuple[Tensor, ...]:
+def plain_tensor_container(obj: AnyTensor) -> tuple[Tensor, ...]:
     """Iteratively flatten an arbitrary output of Tensors into a tuple of Tensors.
 
     Args:

{hyperonnx-1.0.0 → hyperonnx-1.0.2}/hyperonnx/function_rewriter.py

@@ -1,5 +1,5 @@
 """
-Copyright (C) 2025 The HYPERONNX Authors.
+Copyright (C) 2026 The HYPERONNX Authors.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -17,11 +17,12 @@ limitations under the License.
 # pylint: disable=arguments-differ
 
 from collections import defaultdict
+from collections.abc import Sequence
 from contextlib import suppress
 from copy import deepcopy
 from hashlib import md5
 from itertools import chain
-from typing import Dict, List, Literal, Sequence, Set, Tuple, cast
+from typing import Literal, cast
 
 import networkx as nx
 import onnx
@@ -51,7 +52,7 @@ class ComposeOnnxAsFunctionRewriter(Rewriter):
 
     def __init__(self, domain: str, module_spec: Sequence[ModuleSpec]):
         super().__init__(SingleNodePattern().with_domain(domain))
-        self._specs: Dict[str, ModuleSpec] = {
+        self._specs: dict[str, ModuleSpec] = {
             spec["type_name"]: spec for spec in module_spec
         }
 
@@ -122,8 +123,8 @@ class ComposeOnnxAsFunctionRewriter(Rewriter):
         # find the actual input with edges in the graph
         # and remove dangled inputs
         debugable_input_names = spec["input_names"]
-        dangled_input: List[str] = []
-        dangled_or_constant_unsafe: List[str] = []
+        dangled_input: list[str] = []
+        dangled_or_constant_unsafe: list[str] = []
         if len(node.input) != len(debugable_input_names):
             warning(
                 f"Number of inputs on {node.name} is different from spec. "
@@ -143,10 +144,7 @@ class ComposeOnnxAsFunctionRewriter(Rewriter):
                 debug(f"input unconnected to graph node: {input_name}")
                 dangled_or_constant_unsafe.append(input_name)
         if len(dangled_input) < unused:
-            error(
-                f"Expect {unused} unused inputs but "
-                f"enumerated {len(dangled_input)}."
-            )
+            error(f"Expect {unused} unused inputs but enumerated {len(dangled_input)}.")
             reminder = unused - len(dangled_input)
             dangled_input.extend(dangled_or_constant_unsafe[:reminder])
         for j in dangled_input:
@@ -159,8 +157,8 @@ class ComposeOnnxAsFunctionRewriter(Rewriter):
     def _remove_unused_outputs(
         self, graph: OnnxGraph, unused: int, node: NodeProto, spec: ModuleSpec
     ):
-        dangled_output: List[str] = []
-        dangled_output_unsafe: List[str] = []
+        dangled_output: list[str] = []
+        dangled_output_unsafe: list[str] = []
         for i, output_name in enumerate(node.output):
             if output_name in spec["unused_outputs"]:
                 dangled_output.append(output_name)
@@ -170,8 +168,7 @@ class ComposeOnnxAsFunctionRewriter(Rewriter):
                 dangled_output_unsafe.append(output_name)
         if len(dangled_output) < unused:
             error(
-                f"Expect {unused} unused outputs but "
-                f"enumerated {len(dangled_output)}."
+                f"Expect {unused} unused outputs but enumerated {len(dangled_output)}."
             )
             reminder = unused - len(dangled_output)
             # try to drop `reminder` outputs from the last one
@@ -182,7 +179,7 @@ class ComposeOnnxAsFunctionRewriter(Rewriter):
             if j in graph.outputs:
                 graph.remove_output(j)
 
-    def rewrite(self, graph: OnnxGraph, nodes: List[NodeProto]):
+    def rewrite(self, graph: OnnxGraph, nodes: list[NodeProto]):
         node = nodes[0]
         if node.op_type not in self._specs:
             raise KeyError(f"node spec is not found for {node.name}")
@@ -253,7 +250,7 @@ class ComposeNodesToFunctionsRewriter(Rewriter):
         self.register_pre_hook(self._mark_isolated_nodes)
         self._name = name
         self._max_level = 0
-        self._level_processed: Set[int] = set()
+        self._level_processed: set[int] = set()
 
     def _mark_isolated_nodes(self, graph: OnnxGraph) -> OnnxGraph:
         # make a shallow copy of the graph
@@ -269,7 +266,7 @@ class ComposeNodesToFunctionsRewriter(Rewriter):
             self._max_level = max(self._max_level, k + 1)
         return graph
 
-    def _detect_cycles_in_graph(self, graph: OnnxGraph) -> List[List[str]]:
+    def _detect_cycles_in_graph(self, graph: OnnxGraph) -> list[list[str]]:
         """Comprehensive cycle detection in the entire graph."""
         try:
             cycles = list(nx.simple_cycles(graph))
@@ -278,7 +275,7 @@ class ComposeNodesToFunctionsRewriter(Rewriter):
             return []
 
     def _would_create_cycle(
-        self, graph: OnnxGraph, new_node: NodeProto, nodes_to_remove: List[NodeProto]
+        self, graph: OnnxGraph, new_node: NodeProto, nodes_to_remove: list[NodeProto]
     ) -> bool:
         """Check if adding a new node and removing old nodes would create a cycle."""
         # Create a temporary graph by copying the original and removing nodes
@@ -307,7 +304,7 @@ class ComposeNodesToFunctionsRewriter(Rewriter):
             return True  # Assume cycle exists for safety
 
     def _compose_subgraph(
-        self, graph: OnnxGraph, same_level_nodes: List[NodeProto], level: int
+        self, graph: OnnxGraph, same_level_nodes: list[NodeProto], level: int
     ) -> None:
         if len(same_level_nodes) == 0:
             return  # skip empty subgraph
@@ -388,7 +385,7 @@ class ComposeNodesToFunctionsRewriter(Rewriter):
             if direction == "upstream" and i.op_type == "Constant":
                 continue
 
-            same_level_nodes: List[NodeProto] = [i]
+            same_level_nodes: list[NodeProto] = [i]
             # Use BFS to collect neighbors to avoid recursive dependency issues
             visited = {i.name}
             queue = [i.name]
@@ -419,7 +416,7 @@ class ComposeNodesToFunctionsRewriter(Rewriter):
         """Compose downstream nodes into functions."""
         self._compose_neighbors(graph, node, "downstream")
 
-    def rewrite(self, graph: OnnxGraph, nodes: List[NodeProto]):
+    def rewrite(self, graph: OnnxGraph, nodes: list[NodeProto]):
         node = nodes[0]
 
         # Check for cycles before starting composition
@@ -474,10 +471,10 @@ class FuseConstantsToFunctionRewriter(Rewriter):
 
     def __init__(self):
         super().__init__(SingleNodePattern().with_domain(HYPER_DOMAIN))
-        self.users: Dict[str, int] = defaultdict(int)
+        self.users: dict[str, int] = defaultdict(int)
         self.register_pre_hook(self._collect_users)
 
-    def rewrite(self, graph: OnnxGraph, nodes: List[NodeProto]):
+    def rewrite(self, graph: OnnxGraph, nodes: list[NodeProto]):
         node = nodes[0]
         func = graph.functions.get(node.op_type)
         if func is None:
@@ -491,7 +488,7 @@ class FuseConstantsToFunctionRewriter(Rewriter):
             node.op_type = func.name
             self.users[func.name] += 1
 
-        consts: List[Tuple[int, NodeProto]] = []
+        consts: list[tuple[int, NodeProto]] = []
         for i, cst in enumerate(self.get_input_nodes(node)):
             if cst is not None and cst.op_type == "Constant":
                 consts.append((i, cst))
@@ -520,7 +517,7 @@ class EraseOutputTypesRewriter(Rewriter):
         super().__init__(pattern=SingleNodePattern().with_name("__NO_MATCH__"))
         self.register_pre_hook(self._clear_types)
         self.register_post_hook(self._infer_types)
-        self.output_types: List[onnx.TypeProto] = []
+        self.output_types: list[onnx.TypeProto] = []
 
     def _clear_types(self, graph: OnnxGraph):
         for i in graph.output:

{hyperonnx-1.0.0 → hyperonnx-1.0.2}/hyperonnx/hyper_export.py

@@ -14,6 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+from collections.abc import Collection, Container
 from contextlib import suppress
 from inspect import signature
 from io import BytesIO
@@ -21,7 +22,7 @@ from logging import Logger
 from os import PathLike
 from pathlib import Path
 from tempfile import TemporaryDirectory
-from typing import Any, Collection, Container, Dict, List, Optional, Tuple
+from typing import Any
 
 import onnx
 from onnxifier import ONNXIFIER_OPSET, OnnxGraph, PassManager
@@ -48,8 +49,8 @@ from .typing import (
 from .utils import HYPER_DOMAIN, OPTIMIZER_PASSES
 
 
-def _get_input_names(spec: ModuleSpec) -> List[str]:
-    names: List[str] = []
+def _get_input_names(spec: ModuleSpec) -> list[str]:
+    names: list[str] = []
     params = spec["signature"].parameters
 
     def _arg_to_name(name: str, args):
@@ -70,13 +71,13 @@ def _get_input_names(spec: ModuleSpec) -> List[str]:
     return names
 
 
-def _get_output_names(spec: ModuleSpec) -> None | List[str]:
+def _get_output_names(spec: ModuleSpec) -> None | list[str]:
     output = spec.get("output")
     if output is None:
         return None
 
-    def _extract_names_from_dict(out_dict: dict) -> List[str]:
-        names: List[str] = []
+    def _extract_names_from_dict(out_dict: dict) -> list[str]:
+        names: list[str] = []
         for key, value in out_dict.items():
             if isinstance(value, dict):
                 names.extend(_extract_names_from_dict(value))
@@ -99,7 +100,7 @@ def _get_output_names(spec: ModuleSpec) -> None | List[str]:
 
 
 def make_hierarchical_hook(
-    hiera: Container[type[Module]], module_spec: Dict[Module, ModuleSpec], index: int
+    hiera: Container[type[Module]], module_spec: dict[Module, ModuleSpec], index: int
 ) -> HookCallback:
     """Make a forward hook to record spec of modules in `hiera`.
 
@@ -115,8 +116,8 @@ def make_hierarchical_hook(
 
     def _hook(
         module: Module,
-        args: Tuple[Tensor],
-        kwargs: Dict[str, AnyTensor],
+        args: tuple[Tensor],
+        kwargs: dict[str, AnyTensor],
         output: AnyTensor,
     ) -> None:
         spec = module_spec[module]
@@ -140,12 +141,12 @@ def make_hierarchical_hook(
 def trace_module_spec(
     model: Module,
     input_args: tuple,
-    kwargs: Optional[Dict[str, AnyTensor]],
+    kwargs: dict[str, AnyTensor] | None,
     opset_version: int,
     hiera: Container[type[Module]],
-    module_spec: Dict[Module, ModuleSpec],
+    module_spec: dict[Module, ModuleSpec],
     dynamo: bool = False,
-) -> Dict[Module, ModuleSpec]:
+) -> dict[Module, ModuleSpec]:
     """Register forward hooks to modules in the `hiera` and record forward
     information to `module_spec`.
 
@@ -208,13 +209,13 @@ def _export_hiera(
     dynamo: bool,
     external_data: bool,
     do_optimization: bool,
-    external_directory: Optional[str | PathLike],
-    module_spec: Dict[Module, ModuleSpec],
+    external_directory: str | PathLike | None,
+    module_spec: dict[Module, ModuleSpec],
     hiera: Collection[type[Module]],
     logger: Logger,
 ):
-    def _get_sub_spec(module: Module, spec: Dict[Module, ModuleSpec]):
-        child_spec: Dict[Module, ModuleSpec] = default_module_spec()
+    def _get_sub_spec(module: Module, spec: dict[Module, ModuleSpec]):
+        child_spec: dict[Module, ModuleSpec] = default_module_spec()
         for child in module.modules():
             if child in spec:
                 child_spec[child] = spec[child]
@@ -308,16 +309,16 @@ def export_hyper_onnx(  # noqa: C901
     input_args: tuple,
     f: str | PathLike | BytesIO,
     *,
-    kwargs: Optional[Dict[str, AnyTensor]] = None,
-    input_names: Optional[List[str]] = None,
-    output_names: Optional[List[str]] = None,
+    kwargs: dict[str, AnyTensor] | None = None,
+    input_names: list[str] | None = None,
+    output_names: list[str] | None = None,
     opset_version: int = ONNXIFIER_OPSET.version,
     dynamo: bool = False,
     external_data: bool = False,
-    hiera: Optional[Collection[type[Module]]] = None,
-    module_spec: Optional[Dict[Module, ModuleSpec]] = None,
+    hiera: Collection[type[Module]] | None = None,
+    module_spec: dict[Module, ModuleSpec] | None = None,
     do_optimization: bool = True,
-    external_directory: Optional[str | PathLike] = None,
+    external_directory: str | PathLike | None = None,
     **_: Any,  # ignored options
 ) -> Any | None:
     r"""Export a Pytorch module to ONNX format with hierarchical structure,
@@ -450,12 +451,12 @@ def export_hyper_onnx(  # noqa: C901
         else:
             onnx_model = onnx.load_model_from_string(model_path.getvalue())
 
-    typenames: List[str] = [module_spec[i]["type_name"] for i in module_spec]
+    typenames: list[str] = [module_spec[i]["type_name"] for i in module_spec]
     for node in onnx_model.graph.node:
         if node.op_type in typenames:
             node.domain = HYPER_DOMAIN  # add domain tag to run rewriter
     graph = OnnxGraph(onnx_model)
-    passes: List[Any]
+    passes: list[Any]
     passes = [ComposeOnnxAsFunctionRewriter(HYPER_DOMAIN, tuple(module_spec.values()))]
     if do_optimization:
         passes.extend(OPTIMIZER_PASSES)

{hyperonnx-1.0.0 → hyperonnx-1.0.2}/hyperonnx/torch_export.py

@@ -15,10 +15,10 @@ limitations under the License.
 """
 
 import warnings
-from collections.abc import Callable
+from collections.abc import Callable, Mapping, Sequence
 from inspect import signature
 from os import PathLike
-from typing import Any, Mapping, Sequence
+from typing import Any
 
 import torch
 from torch.onnx import OperatorExportTypes

hyperonnx-1.0.2/hyperonnx/transformers/__init__.py

@@ -0,0 +1,10 @@
+from .attention import attention_interface, register_attention_opsets
+from .cache import StaticCache
+from .patch import patch_transformers
+
+__all__ = [
+    "attention_interface",
+    "register_attention_opsets",
+    "StaticCache",
+    "patch_transformers",
+]