Skip to content

Commit 2b487d7

Browse files
justinchubyjustinchuby
committed
Migrate code
1 parent f39a824 commit 2b487d7

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

57 files changed

+15910
-3
lines changed

‎src/onnx_ir/README.md‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
# ONNX IR
2+
3+
An in-memory IR that supports the full ONNX spec, designed for graph construction, analysis and transformation.
4+
5+
## Features ✨
6+
7+
- Full ONNX spec support: all valid models representable by ONNX protobuf, and a subset of invalid models (so you can load and fix them).
8+
- Low memory footprint: mmap'ed external tensors; unified interface for ONNX TensorProto, Numpy arrays and PyTorch Tensors etc. No tensor size limitation. Zero copies.
9+
- Straightforward access patterns: Access value information and traverse the graph topology at ease.
10+
- Robust mutation: Create as many iterators as you like on the graph while mutating it.
11+
- Speed: Performant graph manipulation, serialization/deserialization to Protobuf.
12+
- Pythonic and familiar APIs: Classes define Pythonic apis and still map to ONNX protobuf concepts in an intuitive way.
13+
- No protobuf dependency: The IR does not require protobuf once the model is converted to the IR representation, decoupling from the serialization format.
14+
15+
## Code Organization 🗺️
16+
17+
- [`_protocols.py`](_protocols.py): Interfaces defined for all entities in the IR.
18+
- [`_core.py`](_core.py): Implementation of the core entities in the IR, including `Model`, `Graph`, `Node`, `Value`, and others.
19+
- [`_enums.py`](_enums.py): Definition of the type enums that correspond to the `DataType` and `AttributeType` in `onnx.proto`.
20+
- [`_name_authority.py`](_name_authority.py): The authority for giving names to entities in the graph, used internally.
21+
- [`_linked_list.py`](_linked_list.py): The data structure as the node container in the graph that supports robust iteration and mutation. Internal.
22+
- [`_metadata.py`](_metadata.py): Metadata store for all entities in the IR.

‎src/onnx_ir/__init__.py‎

Lines changed: 158 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,158 @@
1-
# Copyright (c) ONNX Project Contributors
2-
#
3-
# SPDX-License-Identifier: Apache-2.0
1+
# Copyright (c) Microsoft Corporation.
2+
# Licensed under the MIT License.
3+
"""In-memory intermediate representation for ONNX graphs."""
4+
5+
__all__ = [
6+
# Modules
7+
"serde",
8+
"traversal",
9+
"convenience",
10+
"external_data",
11+
"tape",
12+
# IR classes
13+
"Tensor",
14+
"ExternalTensor",
15+
"StringTensor",
16+
"LazyTensor",
17+
"SymbolicDim",
18+
"Shape",
19+
"TensorType",
20+
"OptionalType",
21+
"SequenceType",
22+
"SparseTensorType",
23+
"TypeAndShape",
24+
"Value",
25+
"Attr",
26+
"RefAttr",
27+
"Node",
28+
"Function",
29+
"Graph",
30+
"GraphView",
31+
"Model",
32+
# Constructors
33+
"AttrFloat32",
34+
"AttrFloat32s",
35+
"AttrGraph",
36+
"AttrGraphs",
37+
"AttrInt64",
38+
"AttrInt64s",
39+
"AttrSparseTensor",
40+
"AttrSparseTensors",
41+
"AttrString",
42+
"AttrStrings",
43+
"AttrTensor",
44+
"AttrTensors",
45+
"AttrTypeProto",
46+
"AttrTypeProtos",
47+
"Input",
48+
# Protocols
49+
"ArrayCompatible",
50+
"DLPackCompatible",
51+
"TensorProtocol",
52+
"ValueProtocol",
53+
"ModelProtocol",
54+
"NodeProtocol",
55+
"GraphProtocol",
56+
"GraphViewProtocol",
57+
"AttributeProtocol",
58+
"ReferenceAttributeProtocol",
59+
"SparseTensorProtocol",
60+
"SymbolicDimProtocol",
61+
"ShapeProtocol",
62+
"TypeProtocol",
63+
"MapTypeProtocol",
64+
"FunctionProtocol",
65+
# Enums
66+
"AttributeType",
67+
"DataType",
68+
# Types
69+
"OperatorIdentifier",
70+
# Protobuf compatible types
71+
"TensorProtoTensor",
72+
# Conversion functions
73+
"from_proto",
74+
"from_onnx_text",
75+
"to_proto",
76+
# Convenience constructors
77+
"tensor",
78+
"node",
79+
# Pass infrastructure
80+
"passes",
81+
# IO
82+
"load",
83+
"save",
84+
]
85+
86+
from onnxscript.ir import convenience, external_data, passes, serde, tape, traversal
87+
from onnxscript.ir._convenience._constructors import node, tensor
88+
from onnxscript.ir._core import (
89+
Attr,
90+
AttrFloat32,
91+
AttrFloat32s,
92+
AttrGraph,
93+
AttrGraphs,
94+
AttrInt64,
95+
AttrInt64s,
96+
AttrSparseTensor,
97+
AttrSparseTensors,
98+
AttrString,
99+
AttrStrings,
100+
AttrTensor,
101+
AttrTensors,
102+
AttrTypeProto,
103+
AttrTypeProtos,
104+
ExternalTensor,
105+
Function,
106+
Graph,
107+
GraphView,
108+
Input,
109+
LazyTensor,
110+
Model,
111+
Node,
112+
OptionalType,
113+
RefAttr,
114+
SequenceType,
115+
Shape,
116+
SparseTensorType,
117+
StringTensor,
118+
SymbolicDim,
119+
Tensor,
120+
TensorType,
121+
TypeAndShape,
122+
Value,
123+
)
124+
from onnxscript.ir._enums import (
125+
AttributeType,
126+
DataType,
127+
)
128+
from onnxscript.ir._io import load, save
129+
from onnxscript.ir._protocols import (
130+
ArrayCompatible,
131+
AttributeProtocol,
132+
DLPackCompatible,
133+
FunctionProtocol,
134+
GraphProtocol,
135+
GraphViewProtocol,
136+
MapTypeProtocol,
137+
ModelProtocol,
138+
NodeProtocol,
139+
OperatorIdentifier,
140+
ReferenceAttributeProtocol,
141+
ShapeProtocol,
142+
SparseTensorProtocol,
143+
SymbolicDimProtocol,
144+
TensorProtocol,
145+
TypeProtocol,
146+
ValueProtocol,
147+
)
148+
from onnxscript.ir.serde import TensorProtoTensor, from_onnx_text, from_proto, to_proto
149+
150+
151+
def __set_module() -> None:
152+
"""Set the module of all functions in this module to this public module."""
153+
global_dict = globals()
154+
for name in __all__:
155+
global_dict[name].__module__ = __name__
156+
157+
158+
__set_module()

0 commit comments

Comments
 (0)