burn-import: add some tests for ConstantNode

crates/burn-import/onnx-tests/build.rs

@@ -16,6 +16,10 @@ fn main() {
         .input("tests/clip/clip_opset16.onnx")
         .input("tests/clip/clip_opset7.onnx")
         .input("tests/concat/concat.onnx")
+        .input("tests/constant/constant_f32.onnx")
+        .input("tests/constant/constant_f64.onnx")
+        .input("tests/constant/constant_i32.onnx")
+        .input("tests/constant/constant_i64.onnx")
         .input("tests/constant_of_shape/constant_of_shape.onnx")
         .input("tests/constant_of_shape/constant_of_shape_full_like.onnx")
         .input("tests/conv1d/conv1d.onnx")

crates/burn-import/onnx-tests/tests/constant/constant.py

@@ -0,0 +1,63 @@
+#!/usr/bin/env python3
+
+import torch
+import torch.nn as nn
+
+CONST_VALUE = 2
+
+
+class ConstantModel(nn.Module):
+    def __init__(self, const_dtype: torch.dtype):
+        super().__init__()
+        self.const = torch.tensor(CONST_VALUE).to(const_dtype)
+
+    def forward(self, x):
+        return x + self.const
+
+
+def export_model(model: ConstantModel, dummy_input: torch.Tensor, file_name: str):
+    model.eval()
+    torch.onnx.export(
+        model,
+        dummy_input,
+        file_name,
+        verbose=False,
+        opset_version=16,
+        do_constant_folding=False,
+    )
+    print(f"Finished exporting model to {file_name}")
+
+    # Output some test data for demonstration
+    test_input = dummy_input.clone()
+    print(dummy_input.dtype, "test input:", test_input)
+    output = model.forward(test_input)
+    print(dummy_input.dtype, "test output:", output)
+    print("")
+
+
+def main():
+    device = torch.device("cpu")
+    shape = (2, 3, 4)
+
+    model_f32 = ConstantModel(torch.float32)
+    f32_input = torch.randn(shape, dtype=torch.float32, device=device)
+    export_model(model_f32, f32_input, "constant_f32.onnx")
+
+    model_f64 = ConstantModel(torch.float64)
+    f64_input = torch.randn(shape, dtype=torch.float64, device=device)
+    export_model(model_f64, f64_input, "constant_f64.onnx")
+
+    model_i32 = ConstantModel(torch.int32)
+    i32_input = torch.randint(
+        low=-10, high=10, size=shape, device=device, dtype=torch.int32
+    )
+    export_model(model_i32, i32_input, "constant_i32.onnx")
+
+    model_i64 = ConstantModel(torch.int64)
+    i64_input = torch.randint(
+        low=-10, high=10, size=shape, device=device, dtype=torch.int64
+    )
+    export_model(model_i64, i64_input, "constant_i64.onnx")
+
+if __name__ == "__main__":
+    main()

crates/burn-import/onnx-tests/tests/test_onnx.rs

@@ -25,6 +25,10 @@ include_models!(
     clip_opset16,
     clip_opset7,
     concat,
+    constant_f32,
+    constant_f64,
+    constant_i32,
+    constant_i64,
     constant_of_shape,
     constant_of_shape_full_like,
     conv1d,
@@ -2183,6 +2187,46 @@ mod tests {
         assert_eq!(expected_shape, output.shape());
     }
 
+    #[test]
+    fn add_constant_f32() {
+        let device = Default::default();
+        let model = constant_f32::Model::<Backend>::new(&device);
+        let input = TensorData::zeros::<f32, _>(Shape::from([2, 3, 4]));
+        let expected_output = TensorData::full::<f32, _>(Shape::from([2, 3, 4]), 2f32);
+        let output = model.forward(input.into());
+        assert_eq!(expected_output, output.to_data());
+    }
+
+    #[test]
+    fn add_constant_f64() {
+        let device = Default::default();
+        let model = constant_f64::Model::<Backend>::new(&device);
+        let input = TensorData::zeros::<f64, _>(Shape::from([2, 3, 4]));
+        let expected_output = TensorData::full(Shape::from([2, 3, 4]), 2f32);
+        let output = model.forward(input.into());
+        assert_eq!(expected_output, output.to_data());
+    }
+
+    #[test]
+    fn add_constant_i32() {
+        let device = Default::default();
+        let model = constant_i32::Model::<Backend>::new(&device);
+        let input = TensorData::zeros::<i32, _>(Shape::from([2, 3, 4]));
+        let expected_output = TensorData::full(Shape::from([2, 3, 4]), 2i64);
+        let output = model.forward(input.into());
+        assert_eq!(expected_output, output.to_data());
+    }
+
+    #[test]
+    fn add_constant_i64() {
+        let device = Default::default();
+        let model = constant_i64::Model::<Backend>::new(&device);
+        let input = TensorData::zeros::<i64, _>(Shape::from([2, 3, 4]));
+        let expected_output = TensorData::full(Shape::from([2, 3, 4]), 2i64);
+        let output = model.forward(input.into());
+        assert_eq!(expected_output, output.to_data());
+    }
+
     #[test]
     fn constant_of_shape() {
         // This tests shape is being passed directly to the model