Separating ONNX parsing from burn-import (#1921)

* separating onnx parsing from burn-import

* ran clippy and cargo-fmt

* removed unused deps from onnx-ir

* fixed clippy warnings that were causing run-checks to fail

* removed dead code

* removed unused dependencies from burn-import

* updated contributor-book, updated publish.yml, added readme

* update cargo lock

* formatted md document with prettier, rephrased sentence

* missed the errors with reduce_prod_conversion during merge

* formatted onnx-to-burn-conversion-tool.md, forgot to save

.github/workflows/publish.yml

@@ -161,3 +161,9 @@ jobs:
     with:
       crate: burn-import
     secrets: inherit
+
+  publish-onnx-ir:
+    uses: tracel-ai/burn/.github/workflows/publish-template.yml@main
+    with:
+      crate: onnx-ir
+    secrets: inherit

contributor-book/src/guides/adding-a-new-operation-to-burn.md

@@ -117,7 +117,8 @@ plug in your operator in terms of \\(x\\) and \\(y\\), and just swap out the var
 
 ### Testing autodiff
 
-For testing the `autodiff` operations, please refer to [this section](../getting-started/testing.md).
+For testing the `autodiff` operations, please refer to
+[this section](../getting-started/testing.md).
 
 ## Adding the Op to other backends
 
@@ -199,11 +200,15 @@ Generating the ONNX test files or tests is already covered
 about the specific changes you need to make when adding new operators after you have generated the
 tests.
 
-The crate is divided into two sections `src/burn` and `src/onnx`. The code under the former
-corresponds to the operation you've implemented earlier in this guide, and the latter to the
-operations defined in the ONNX specification. So when you are loading a model, the operator is first
-parsed to an intermediate representation defined by `src/onnx`, and then mapped to a Burn operation
-defined under `src/burn/node`.
+Changes will need to be made to both `onnx-ir` and `burn-import`. The code within `onnx-ir` defines
+how to parse the nodes in an onnx file and produces the intermediate representation. The code within
+`burn-import` is divided into two sections: `src/onnx` and `src/burn`. The code under the former
+maps that intermediate representation to one used for code generation and the latter defines how to
+generate code for the operator you've implemented earlier in this guide.
+
+So when you are loading a model, the operator is first parsed to an intermediate representation
+defined by `burn-import` and then mapped to a Burn operation defined under `src/burn/node`; the
+mapping from onnx to burn is aptly defined in `src/onnx/to_burn`
 
 Let's review the changes made for powf starting from `src/burn` and moving to `src/onnx`:
 
@@ -218,17 +223,20 @@ Let's review the changes made for powf starting from `src/burn` and moving to `s
    [`{op}_conversion` function](https://github.com/tracel-ai/burn/blob/0ee2021567b3725907df5fd1a905ce60b1aca096/crates/burn-import/src/onnx/to_burn.rs#L717)
    that maps the ONNX node to the binary type
 3. Specify how dimensions for the output should be derived in
-   [crates/burn-import/src/onnx/dim_inference.rs](https://github.com/tracel-ai/burn/blob/0ee2021567b3725907df5fd1a905ce60b1aca096/crates/burn-import/src/onnx/dim_inference.rs#L55)
+   [crates/onnx-ir/src/dim_inference.rs](https://github.com/tracel-ai/burn/blob/d4ae82b21ac3dd1def01bd380ab7ea4d3293eccb/crates/onnx-ir/src/dim_inference.rs#L17)
 
 And you're done! Congrats, you just fully added a new operation to burn, and we are all one step
 closer to the answer to [Are we learning yet?](https://www.arewelearningyet.com/) being "Yes, and
 it's freaking fast!". Buy yourself a coffee.
 
-[^supertrait]: for more on supertraits see
+[^supertrait]:
+    for more on supertraits see
     [the advanced trait section of the rust book](https://doc.rust-lang.org/book/ch19-03-advanced-traits.html#using-supertraits-to-require-one-traits-functionality-within-another-trait)
 
-[^autodiff]: wiki link for
+[^autodiff]:
+    wiki link for
     [automatic differentiation](https://en.wikipedia.org/wiki/Automatic_differentiation)
 
-[^absolute_units]: for more information on unit structs see
+[^absolute_units]:
+    for more information on unit structs see
     [the defining and instantiating structs section of the rust book](https://doc.rust-lang.org/book/ch05-01-defining-structs.html#unit-like-structs-without-any-fields)

contributor-book/src/guides/onnx-to-burn-conversion-tool.md

@@ -16,7 +16,17 @@ For an introduction to ONNX import in Burn, see
     - [Design Goals](#design-goals)
     - [Design Decisions](#design-decisions)
   - [Adding New Operators](#adding-new-operators)
-    - [Implementing a New Operator](#implementing-a-new-operator)
+  - [Implementing a New Operator](#implementing-a-new-operator)
+    - [Step 1: Visibility](#step-1-visibility)
+    - [Step 2: Node Implementation](#step-2-node-implementation)
+      - [Within Onnx-IR](#within-onnx-ir)
+      - [Within burn-import](#within-burn-import)
+    - [Step 3: Registering New Operations](#step-3-registering-new-operations)
+    - [Step 4: Create a Config Function](#step-4-create-a-config-function)
+    - [Step 5: Dimension Inference](#step-5-dimension-inference)
+    - [Step 6: Integrate into the Graph Building Process](#step-6-integrate-into-the-graph-building-process)
+    - [Step 7: Add Newly Supported Op!](#step-7-add-newly-supported-op)
+    - [Misc:](#misc)
   - [Testing](#testing)
   - [Resources](#resources)
 
@@ -91,6 +101,22 @@ located in the `src/burn/node/` directory.
 
 ### Step 2: Node Implementation
 
+#### Within Onnx-IR
+
+If the node type does not exist within the
+[`NodeType` enum](https://github.com/tracel-ai/burn/blob/d4ae82b21ac3dd1def01bd380ab7ea4d3293eccb/crates/onnx-ir/src/ir.rs#L246),
+it will need to be added (support for custom operators is planned). If the node might be provided an
+input which is a constant or the output of an identity node, it will need to be added to the list of
+nodeTypes
+[checked for constants](https://github.com/tracel-ai/burn/blob/d4ae82b21ac3dd1def01bd380ab7ea4d3293eccb/crates/onnx-ir/src/from_onnx.rs#L21).
+The node will need to be added to `dim_inference`, and in most cases the work parsing side will be
+done. If a node requires extra parsing (such as handling an edge case like potentially remapping an
+unsqueeze to a reshape) the best place for that is after check constants and prior to dim_inference
+in
+[`OnnxGraphBuilder::Build`](https://github.com/tracel-ai/burn/blob/d4ae82b21ac3dd1def01bd380ab7ea4d3293eccb/crates/onnx-ir/src/from_onnx.rs#L221)
+
+#### Within burn-import
+
 Create a new file named `<operation_name>.rs` in the `src/burn/node/` directory.  
 This file will define the structure and functionality of your new operation. By convention, the
 necessary information for carrying out an operation is encapsulated within a struct named

crates/burn-import/Cargo.toml

@@ -20,30 +20,23 @@ pytorch = ["burn/record-item-custom-serde", "thiserror", "zip"]
 
 [dependencies]
 burn = { path = "../burn", version = "0.14.0", features = ["ndarray"] }
-
-bytemuck = { workspace = true }
+onnx-ir = { path = "../onnx-ir" }
 candle-core = { workspace = true }
 derive-new = { workspace = true }
 half = { workspace = true }
 log = { workspace = true }
 proc-macro2 = { workspace = true }
-protobuf = { workspace = true, features = ["with-bytes"] }
 quote = { workspace = true }
 regex = { workspace = true }
 rust-format = { workspace = true, features = ["token_stream", "post_process"] }
 serde = { workspace = true, features = ["derive"] }
 serde_json = { workspace = true, features = ["std"] }
-strum = { workspace = true }
-strum_macros = { workspace = true }
 syn = { workspace = true, features = ["parsing"] }
 thiserror = { workspace = true, optional = true }
 tracing-core = { workspace = true }
 tracing-subscriber = { workspace = true }
 zip = { workspace = true, optional = true }
 
-[build-dependencies]
-protobuf-codegen = { workspace = true }
-
 [dev-dependencies]
 pretty_assertions = { workspace = true }
 rstest = { workspace = true }

crates/burn-import/src/onnx/mod.rs

@@ -1,14 +1,3 @@
-mod coalesce;
-mod dim_inference;
-mod from_onnx;
-mod ir;
-mod node_remap;
 mod op_configuration;
-mod proto_conversion;
-mod protos;
 mod to_burn;
-
 pub use to_burn::*;
-
-pub use from_onnx::parse_onnx;
-pub use ir::OnnxGraph;

crates/burn-import/src/onnx/op_configuration.rs

@@ -5,8 +5,8 @@ use burn::nn::{
     PaddingConfig2d,
 };
 
-use super::ir::{ArgType, AttributeValue, Data, Node};
 use crate::burn::node::resize::ResizeMode;
+use onnx_ir::ir::{ArgType, AttributeValue, Data, Node};
 
 /// Create a Conv1dConfig from the attributes of the node
 pub fn conv1d_config(curr: &Node) -> Conv1dConfig {