├── .github └── workflows │ └── publish-to-pypi.yml ├── .gitignore ├── LICENSE ├── README.md ├── __init__.py ├── examples ├── custom_tracer.py ├── invert.py ├── model_tracer.py ├── primitive_library.py ├── proxy_based_graph_creation.py ├── replace_op.py ├── subgraph_rewriter_basic_use.py └── wrap_outpt_dynamically.py ├── onefx ├── __init__.py ├── _compatibility.py ├── _pytree.py ├── _symbolic_trace.py ├── annotate.py ├── graph.py ├── graph_module.py ├── immutable_collections.py ├── interpreter.py ├── node.py ├── operator_schemas.py ├── passes │ ├── __init__.py │ ├── annotate_getitem_nodes.py │ ├── backends │ │ ├── __init__.py │ │ └── cudagraphs.py │ ├── dialect │ │ ├── __init__.py │ │ └── common │ │ │ ├── __init__.py │ │ │ └── cse_pass.py │ ├── fake_tensor_prop.py │ ├── graph_drawer.py │ ├── graph_manipulation.py │ ├── infra │ │ ├── __init__.py │ │ ├── partitioner.py │ │ ├── pass_base.py │ │ └── pass_manager.py │ ├── net_min_base.py │ ├── operator_support.py │ ├── param_fetch.py │ ├── pass_manager.py │ ├── reinplace.py │ ├── shape_prop.py │ ├── split_module.py │ ├── split_utils.py │ ├── splitter_base.py │ ├── tests │ │ ├── __init__.py │ │ └── test_pass_manager.py │ ├── tools_common.py │ └── utils │ │ ├── __init__.py │ │ ├── common.py │ │ ├── fuser_utils.py │ │ └── matcher_utils.py ├── proxy.py ├── subgraph_rewriter.py ├── tensor_type.py ├── traceback.py ├── utils │ ├── __init__.py │ ├── _mangling.py │ ├── _pytree.py │ ├── _storage_weak_ref.py │ ├── _submodule.py │ ├── importer.py │ ├── package_exporter.py │ └── package_importer.py └── version.py ├── setup.py └── test ├── __init__.py ├── modules ├── __init__.py └── basic.py ├── simple_test.py ├── test_alexnet.py ├── test_conv_bn_fuse.py ├── test_cpu_performance_profiler.py ├── test_dce_pass.py ├── test_future.py ├── test_fx_feature_extraction.py ├── test_fx_param_shape_control_flow.py ├── test_global_wrap.py ├── test_pass_infra.py └── utils ├── __init__.py └── basic.py /.github/workflows/publish-to-pypi.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/.github/workflows/publish-to-pypi.yml -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/.gitignore -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/LICENSE -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/README.md -------------------------------------------------------------------------------- /__init__.py: -------------------------------------------------------------------------------- 1 | from . import onefx 2 | -------------------------------------------------------------------------------- /examples/custom_tracer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/examples/custom_tracer.py -------------------------------------------------------------------------------- /examples/invert.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/examples/invert.py -------------------------------------------------------------------------------- /examples/model_tracer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/examples/model_tracer.py -------------------------------------------------------------------------------- /examples/primitive_library.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/examples/primitive_library.py -------------------------------------------------------------------------------- /examples/proxy_based_graph_creation.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/examples/proxy_based_graph_creation.py -------------------------------------------------------------------------------- /examples/replace_op.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/examples/replace_op.py -------------------------------------------------------------------------------- /examples/subgraph_rewriter_basic_use.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/examples/subgraph_rewriter_basic_use.py -------------------------------------------------------------------------------- /examples/wrap_outpt_dynamically.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/examples/wrap_outpt_dynamically.py -------------------------------------------------------------------------------- /onefx/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/__init__.py -------------------------------------------------------------------------------- /onefx/_compatibility.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/_compatibility.py -------------------------------------------------------------------------------- /onefx/_pytree.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/_pytree.py -------------------------------------------------------------------------------- /onefx/_symbolic_trace.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/_symbolic_trace.py -------------------------------------------------------------------------------- /onefx/annotate.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/annotate.py -------------------------------------------------------------------------------- /onefx/graph.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/graph.py -------------------------------------------------------------------------------- /onefx/graph_module.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/graph_module.py -------------------------------------------------------------------------------- /onefx/immutable_collections.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/immutable_collections.py -------------------------------------------------------------------------------- /onefx/interpreter.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/interpreter.py -------------------------------------------------------------------------------- /onefx/node.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/node.py -------------------------------------------------------------------------------- /onefx/operator_schemas.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /onefx/passes/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/passes/__init__.py -------------------------------------------------------------------------------- /onefx/passes/annotate_getitem_nodes.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/passes/annotate_getitem_nodes.py -------------------------------------------------------------------------------- /onefx/passes/backends/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /onefx/passes/backends/cudagraphs.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/passes/backends/cudagraphs.py -------------------------------------------------------------------------------- /onefx/passes/dialect/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /onefx/passes/dialect/common/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /onefx/passes/dialect/common/cse_pass.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/passes/dialect/common/cse_pass.py -------------------------------------------------------------------------------- /onefx/passes/fake_tensor_prop.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/passes/fake_tensor_prop.py -------------------------------------------------------------------------------- /onefx/passes/graph_drawer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/passes/graph_drawer.py -------------------------------------------------------------------------------- /onefx/passes/graph_manipulation.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/passes/graph_manipulation.py -------------------------------------------------------------------------------- /onefx/passes/infra/__init__.py: -------------------------------------------------------------------------------- 1 | 2 | from . import pass_manager 3 | -------------------------------------------------------------------------------- /onefx/passes/infra/partitioner.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/passes/infra/partitioner.py -------------------------------------------------------------------------------- /onefx/passes/infra/pass_base.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/passes/infra/pass_base.py -------------------------------------------------------------------------------- /onefx/passes/infra/pass_manager.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/passes/infra/pass_manager.py -------------------------------------------------------------------------------- /onefx/passes/net_min_base.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/passes/net_min_base.py -------------------------------------------------------------------------------- /onefx/passes/operator_support.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/passes/operator_support.py -------------------------------------------------------------------------------- /onefx/passes/param_fetch.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/passes/param_fetch.py -------------------------------------------------------------------------------- /onefx/passes/pass_manager.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/passes/pass_manager.py -------------------------------------------------------------------------------- /onefx/passes/reinplace.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/passes/reinplace.py -------------------------------------------------------------------------------- /onefx/passes/shape_prop.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/passes/shape_prop.py -------------------------------------------------------------------------------- /onefx/passes/split_module.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/passes/split_module.py -------------------------------------------------------------------------------- /onefx/passes/split_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/passes/split_utils.py -------------------------------------------------------------------------------- /onefx/passes/splitter_base.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/passes/splitter_base.py -------------------------------------------------------------------------------- /onefx/passes/tests/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /onefx/passes/tests/test_pass_manager.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/passes/tests/test_pass_manager.py -------------------------------------------------------------------------------- /onefx/passes/tools_common.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/passes/tools_common.py -------------------------------------------------------------------------------- /onefx/passes/utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/passes/utils/__init__.py -------------------------------------------------------------------------------- /onefx/passes/utils/common.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/passes/utils/common.py -------------------------------------------------------------------------------- /onefx/passes/utils/fuser_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/passes/utils/fuser_utils.py -------------------------------------------------------------------------------- /onefx/passes/utils/matcher_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/passes/utils/matcher_utils.py -------------------------------------------------------------------------------- /onefx/proxy.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/proxy.py -------------------------------------------------------------------------------- /onefx/subgraph_rewriter.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/subgraph_rewriter.py -------------------------------------------------------------------------------- /onefx/tensor_type.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/tensor_type.py -------------------------------------------------------------------------------- /onefx/traceback.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/traceback.py -------------------------------------------------------------------------------- /onefx/utils/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /onefx/utils/_mangling.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/utils/_mangling.py -------------------------------------------------------------------------------- /onefx/utils/_pytree.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/utils/_pytree.py -------------------------------------------------------------------------------- /onefx/utils/_storage_weak_ref.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/utils/_storage_weak_ref.py -------------------------------------------------------------------------------- /onefx/utils/_submodule.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/utils/_submodule.py -------------------------------------------------------------------------------- /onefx/utils/importer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/utils/importer.py -------------------------------------------------------------------------------- /onefx/utils/package_exporter.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/utils/package_exporter.py -------------------------------------------------------------------------------- /onefx/utils/package_importer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/utils/package_importer.py -------------------------------------------------------------------------------- /onefx/version.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/onefx/version.py -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/setup.py -------------------------------------------------------------------------------- /test/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /test/modules/__init__.py: -------------------------------------------------------------------------------- 1 | from .basic import MyModule -------------------------------------------------------------------------------- /test/modules/basic.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/test/modules/basic.py -------------------------------------------------------------------------------- /test/simple_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/test/simple_test.py -------------------------------------------------------------------------------- /test/test_alexnet.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/test/test_alexnet.py -------------------------------------------------------------------------------- /test/test_conv_bn_fuse.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/test/test_conv_bn_fuse.py -------------------------------------------------------------------------------- /test/test_cpu_performance_profiler.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/test/test_cpu_performance_profiler.py -------------------------------------------------------------------------------- /test/test_dce_pass.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/test/test_dce_pass.py -------------------------------------------------------------------------------- /test/test_future.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/test/test_future.py -------------------------------------------------------------------------------- /test/test_fx_feature_extraction.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/test/test_fx_feature_extraction.py -------------------------------------------------------------------------------- /test/test_fx_param_shape_control_flow.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/test/test_fx_param_shape_control_flow.py -------------------------------------------------------------------------------- /test/test_global_wrap.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/test/test_global_wrap.py -------------------------------------------------------------------------------- /test/test_pass_infra.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/test/test_pass_infra.py -------------------------------------------------------------------------------- /test/utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .basic import add_n -------------------------------------------------------------------------------- /test/utils/basic.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Oneflow-Inc/one-fx/HEAD/test/utils/basic.py --------------------------------------------------------------------------------