feat: added expand to import

tracel-ai · May 31, 2024 · 10cda31 · 10cda31
1 parent da8a522
commit 10cda31
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 8 deletions.
diff --git a/crates/burn-import/onnx-tests/tests/onnx_tests.rs b/crates/burn-import/onnx-tests/tests/onnx_tests.rs
@@ -1124,13 +1124,12 @@ mod tests {
         let device = Default::default();
         let model: expand::Model<Backend> = expand::Model::new(&device);
 
-        let input1 = Tensor::<Backend, 4>::from_floats([[[[-1.0, 1.0, 42.0, 3.0]]]], &device);
-        let input2 = Tensor::<Backend, 1, Int>::from_ints([3, 2], &device);
+        let input1 = Tensor::<Backend, 2>::from_floats([[-1.0], [1.0]], &device);
 
-        // let output = model.forward(input1, input2);
-        // let expected_shape = Shape::from([3, 2]);
-        //
-        // assert_eq!(output.shape(), expected_shape);
+        let output = model.forward(input1);
+        let expected_shape = Shape::from([2, 2]);
+
+        assert_eq!(output.shape(), expected_shape);
     }
 
     #[test]

diff --git a/crates/burn-import/src/onnx/from_onnx.rs b/crates/burn-import/src/onnx/from_onnx.rs
@@ -17,12 +17,13 @@ use super::ir::{ArgType, Argument, Node, NodeType};
 
 use protobuf::Message;
 
-const LIFT_CONSTANTS_FOR_NODE_TYPES: [NodeType; 9] = [
+const LIFT_CONSTANTS_FOR_NODE_TYPES: [NodeType; 10] = [
     NodeType::BatchNormalization,
     NodeType::Clip,
     NodeType::Conv1d,
     NodeType::Conv2d,
     NodeType::Dropout,
+    NodeType::Expand,
     NodeType::Reshape,
     NodeType::Unsqueeze,
     NodeType::ReduceSum,

diff --git a/crates/burn-import/src/onnx/to_burn.rs b/crates/burn-import/src/onnx/to_burn.rs
@@ -878,7 +878,6 @@ impl OnnxGraph {
     fn expand_conversion(node: Node) -> ExpandNode {
         let input = node.inputs.first().unwrap().to_tensor_type();
         let output = node.outputs.first().unwrap().to_tensor_type();
-        println!("{:?}", node);
         let shape = expand_config(&node);
 
         ExpandNode::new(input, output, shape)