feat: onnx ceil & round

Author: tye-singwa

crates/burn-autodiff/src/ops/tensor.rs

@@ -2046,9 +2046,9 @@ impl<B: Backend, C: CheckpointStrategy> FloatTensorOps<Self> for Autodiff<B, C>
         {
             OpsKind::Tracked(preps) => preps.finish(
                 (tensor.primitive.shape(), B::float_device(&tensor.primitive)),
-                B::float_floor(tensor.primitive),
+                B::float_ceil(tensor.primitive),
             ),
-            OpsKind::UnTracked(preps) => preps.finish(B::float_floor(tensor.primitive)),
+            OpsKind::UnTracked(preps) => preps.finish(B::float_ceil(tensor.primitive)),
         }
     }
 

crates/burn-import/SUPPORTED-ONNX-OPS.md

@@ -33,7 +33,7 @@ functionality.
 | [BlackmanWindow][21]             | ❌             | ❌           |
 | [Cast][22]                       | ✅             | ✅           |
 | [CastLike][23]                   | ❌             | ❌           |
-| [Ceil][24]                       | ❌             | ✅           |
+| [Ceil][24]                       | ✅             | ✅           |
 | [Celu][25]                       | ❌             | ❌           |
 | [CenterCropPad][26]              | ❌             | ❌           |
 | [Clip][27]                       | ✅             | ✅           |
@@ -158,7 +158,7 @@ functionality.
 | [ReverseSequence][144]           | ❌             | ❌           |
 | [RNN][145]                       | ❌             | ✅           |
 | [RoiAlign][146]                  | ❌             | ❌           |
-| [Round][147]                     | ❌             | ✅           |
+| [Round][147]                     | ✅             | ✅           |
 | [Scan][148]                      | ❌             | ❌           |
 | [Scatter][149]                   | ❌             | ✅           |
 | [ScatterElements][150]           | ❌             | ❌           |

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

@@ -14,6 +14,7 @@ fn main() {
         .input("tests/avg_pool2d/avg_pool2d.onnx")
         .input("tests/batch_norm/batch_norm.onnx")
         .input("tests/cast/cast.onnx")
+        .input("tests/ceil/ceil.onnx")
         .input("tests/clip/clip.onnx")
         .input("tests/concat/concat.onnx")
         .input("tests/constant/constant_f32.onnx")
@@ -109,6 +110,7 @@ fn main() {
         .input("tests/resize/resize_2d_bicubic_scale.onnx")
         .input("tests/resize/resize_2d_bilinear_scale.onnx")
         .input("tests/resize/resize_2d_nearest_scale.onnx")
+        .input("tests/round/round.onnx")
         .input("tests/shape/shape.onnx")
         .input("tests/sigmoid/sigmoid.onnx")
         .input("tests/sign/sign.onnx")

crates/burn-import/onnx-tests/tests/ceil/ceil.onnx

@@ -0,0 +1,12 @@
+pytorch2.3.1:Y
+
+onnx::Ceil_0
1/Ceil"Ceil
+main_graphZ
+onnx::Ceil_0
+
+

+b
+
1
+
+

+B

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

@@ -0,0 +1,40 @@
+#!/usr/bin/env python3
+
+# Used to generate model: onnx-tests/tests/ceil/ceil.onnx
+
+import torch
+import torch.nn as nn
+import onnx
+
+class CeilModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return torch.ceil(x)
+
+def main():
+    model = CeilModel()
+    model.eval()
+
+    test_input = torch.tensor([-0.5, 1.5, 2.1])
+
+    onnx_file = "ceil.onnx"
+
+    torch.onnx.export(
+        model,
+        test_input,
+        onnx_file,
+        opset_version=16,
+    )
+
+    print(f"Finished exporting model to {onnx_file}")
+    print(f"Test input data of ones: {test_input}")
+    print(f"Test input data shape of ones: {test_input.shape}")
+    output = model.forward(test_input)
+    print(f"Test output data shape: {output.shape}")
+    print(f"Test output: {output}")
+
+
+if __name__ == '__main__':
+    main()

crates/burn-import/onnx-tests/tests/ceil/mod.rs

@@ -0,0 +1,27 @@
+use crate::include_models;
+include_models!(ceil);
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use burn::tensor::{Tensor, ops::FloatElem};
+
+    type Backend = burn_ndarray::NdArray<f32>;
+    type FT = FloatElem<Backend>;
+
+    #[test]
+    fn ceil_test() {
+        // Test for ceil
+        let device = Default::default();
+        let model = ceil::Model::<Backend>::new(&device);
+
+        let input = Tensor::<Backend, 1>::from_floats([-0.5, 1.5, 2.1], &device);
+        let expected = Tensor::<Backend, 1>::from_floats([0., 2., 3.], &device);
+
+        let output = model.forward(input);
+
+        output
+            .to_data()
+            .assert_approx_eq::<FT>(&expected.to_data(), burn::tensor::Tolerance::default());
+    }
+}

crates/burn-import/onnx-tests/tests/round/mod.rs

@@ -0,0 +1,27 @@
+use crate::include_models;
+include_models!(round);
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use burn::tensor::{Tensor, ops::FloatElem};
+
+    type Backend = burn_ndarray::NdArray<f32>;
+    type FT = FloatElem<Backend>;
+
+    #[test]
+    fn round_test() {
+        // Test for round
+        let device = Default::default();
+        let model = round::Model::<Backend>::new(&device);
+
+        let input = Tensor::<Backend, 1>::from_floats([-0.5, 1.5, 2.1], &device);
+        let expected = Tensor::<Backend, 1>::from_floats([0., 2., 2.], &device);
+
+        let output = model.forward(input);
+
+        output
+            .to_data()
+            .assert_approx_eq::<FT>(&expected.to_data(), burn::tensor::Tolerance::default());
+    }
+}

crates/burn-import/onnx-tests/tests/round/round.onnx

@@ -0,0 +1,12 @@
+pytorch2.3.1:]
+!
+onnx::Round_0
1/Round"Round
+main_graphZ
+onnx::Round_0
+
+

+b
+
1
+
+

+B

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

@@ -0,0 +1,40 @@
+#!/usr/bin/env python3
+
+# Used to generate model: onnx-tests/tests/round/round.onnx
+
+import torch
+import torch.nn as nn
+import onnx
+
+class RoundModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return torch.round(x)
+
+def main():
+    model = RoundModel()
+    model.eval()
+
+    test_input = torch.tensor([-0.5, 1.5, 2.1])
+
+    onnx_file = "round.onnx"
+
+    torch.onnx.export(
+        model,
+        test_input,
+        onnx_file,
+        opset_version=16,
+    )
+
+    print(f"Finished exporting model to {onnx_file}")
+    print(f"Test input data of ones: {test_input}")
+    print(f"Test input data shape of ones: {test_input.shape}")
+    output = model.forward(test_input)
+    print(f"Test output data shape: {output.shape}")
+    print(f"Test output: {output}")
+
+
+if __name__ == '__main__':
+    main()

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

@@ -9,6 +9,7 @@ pub mod argmax;
 pub mod avg_pool;
 pub mod batch_norm;
 pub mod cast;
+pub mod ceil;
 pub mod clip;
 pub mod concat;
 pub mod constant;
@@ -69,6 +70,7 @@ pub mod reduce_sum;
 pub mod relu;
 pub mod reshape;
 pub mod resize;
+pub mod round;
 pub mod shape;
 pub mod sigmoid;
 pub mod sign;

crates/burn-import/src/burn/node/base.rs

@@ -2,8 +2,8 @@ use std::marker::PhantomData;
 
 use super::{
     argmax::ArgMaxNode, avg_pool1d::AvgPool1dNode, avg_pool2d::AvgPool2dNode,
-    batch_norm::BatchNormNode, binary::BinaryNode, clip::ClipNode, concat::ConcatNode,
-    constant::ConstantNode, constant_of_shape::ConstantOfShapeNode,
+    batch_norm::BatchNormNode, binary::BinaryNode, ceil::CeilNode, clip::ClipNode,
+    concat::ConcatNode, constant::ConstantNode, constant_of_shape::ConstantOfShapeNode,
     conv_transpose_1d::ConvTranspose1dNode, conv_transpose_2d::ConvTranspose2dNode,
     conv_transpose_3d::ConvTranspose3dNode, conv1d::Conv1dNode, conv2d::Conv2dNode,
     conv3d::Conv3dNode, dropout::DropoutNode, expand::ExpandNode, floor::FloorNode,
@@ -13,9 +13,9 @@ use super::{
     max_pool2d::MaxPool2dNode, mean::MeanNode, one_hot::OneHotNode, pad::PadNode, prelu::PReluNode,
     random_normal::RandomNormalNode, random_normal_like::RandomNormalLikeNode,
     random_uniform::RandomUniformNode, random_uniform_like::RandomUniformLikeNode,
-    range::RangeNode, reshape::ReshapeNode, resize::ResizeNode, slice::SliceNode, split::SplitNode,
-    squeeze::SqueezeNode, sum::SumNode, tile::TileNode, top_k::TopKNode, trilu::TriluNode,
-    unary::UnaryNode, unsqueeze::UnsqueezeNode,
+    range::RangeNode, reshape::ReshapeNode, resize::ResizeNode, round::RoundNode, slice::SliceNode,
+    split::SplitNode, squeeze::SqueezeNode, sum::SumNode, tile::TileNode, top_k::TopKNode,
+    trilu::TriluNode, unary::UnaryNode, unsqueeze::UnsqueezeNode,
 };
 use crate::burn::{BurnImports, Scope, Type};
 use burn::record::PrecisionSettings;
@@ -102,6 +102,7 @@ pub enum Node<PS: PrecisionSettings> {
     Dropout(DropoutNode),
     Expand(ExpandNode),
     Floor(FloorNode),
+    Ceil(CeilNode),
     Gather(GatherNode),
     GatherElements(GatherElementsNode),
     Gemm(GemmNode),
@@ -118,6 +119,7 @@ pub enum Node<PS: PrecisionSettings> {
     Range(RangeNode),
     Reshape(ReshapeNode),
     Resize(ResizeNode),
+    Round(RoundNode),
     Slice(SliceNode),
     Squeeze(SqueezeNode),
     Split(SplitNode),
@@ -160,6 +162,7 @@ macro_rules! match_all {
             Node::Dropout(node) => $func(node),
             Node::Expand(node) => $func(node),
             Node::Floor(node) => $func(node),
+            Node::Ceil(node) => $func(node),
             Node::Gather(node) => $func(node),
             Node::GatherElements(node) => $func(node),
             Node::Gemm(node) => $func(node),
@@ -176,6 +179,7 @@ macro_rules! match_all {
             Node::Range(node) => $func(node),
             Node::Reshape(node) => $func(node),
             Node::Resize(node) => $func(node),
+            Node::Round(node) => $func(node),
             Node::Slice(node) => $func(node),
             Node::Squeeze(node) => $func(node),
             Node::Sum(node) => $func(node),
@@ -226,6 +230,7 @@ impl<PS: PrecisionSettings> Node<PS> {
             Node::Dropout(_) => "dropout",
             Node::Expand(_) => "expand",
             Node::Floor(_) => "floor",
+            Node::Ceil(_) => "ceil",
             Node::Gather(_) => "gather",
             Node::GatherElements(_) => "gather_elements",
             Node::Gemm(_) => "gemm",
@@ -242,6 +247,7 @@ impl<PS: PrecisionSettings> Node<PS> {
             Node::Range(_) => "range",
             Node::Reshape(_) => "reshape",
             Node::Resize(_) => "resize",
+            Node::Round(_) => "round",
             Node::Slice(_) => "slice",
             Node::Squeeze(_) => "squeeze",
             Node::Sum(_) => "add",

crates/burn-import/src/burn/node/ceil.rs

@@ -0,0 +1,88 @@
+use super::{Node, NodeCodegen};
+use crate::burn::{Scope, TensorType, Type};
+use burn::record::PrecisionSettings;
+use proc_macro2::TokenStream;
+use quote::quote;
+
+#[derive(Debug, Clone, new)]
+pub struct CeilNode {
+    pub input: TensorType,
+    pub output: TensorType,
+}
+
+impl<PS: PrecisionSettings> NodeCodegen<PS> for CeilNode {
+    fn input_types(&self) -> Vec<Type> {
+        vec![Type::Tensor(self.input.clone())]
+    }
+
+    fn output_types(&self) -> Vec<Type> {
+        vec![Type::Tensor(self.output.clone())]
+    }
+
+    fn forward(&self, scope: &mut Scope, node_position: usize) -> TokenStream {
+        let input = scope.tensor_use_owned(&self.input, node_position);
+        let output = &self.output.name;
+
+        quote! {
+            let #output = #input.ceil();
+        }
+    }
+
+    fn into_node(self) -> Node<PS> {
+        Node::Ceil(self)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use burn::record::FullPrecisionSettings;
+
+    use super::*;
+    use crate::burn::{
+        TensorType,
+        graph::BurnGraph,
+        node::{ceil::CeilNode, test::assert_tokens},
+    };
+
+    #[test]
+    fn test_codegen_nodes() {
+        let mut graph = BurnGraph::<FullPrecisionSettings>::default();
+
+        graph.register(CeilNode::new(
+            TensorType::new_float("tensor1", 1),
+            TensorType::new_float("tensor2", 1),
+        ));
+
+        graph.register_input_output(vec!["tensor1".to_string()], vec!["tensor2".to_string()]);
+
+        let expected = quote! {
+            use burn::{
+                module::Module,
+                tensor::{backend::Backend, Tensor},
+            };
+
+            #[derive(Module, Debug)]
+            pub struct Model<B: Backend> {
+                phantom: core::marker::PhantomData<B>,
+                device: burn::module::Ignored<B::Device>,
+            }
+
+            impl<B: Backend> Model<B> {
+                #[allow(unused_variables)]
+                pub fn new(device: &B::Device) -> Self {
+                    Self {
+                        phantom: core::marker::PhantomData,
+                        device: burn::module::Ignored(device.clone()),
+                    }
+                }
+                #[allow(clippy::let_and_return, clippy::approx_constant)]
+                pub fn forward(&self, tensor1: Tensor<B, 1>) -> Tensor<B, 1> {
+                    let tensor2 = tensor1.ceil();
+                    tensor2
+                }
+            }
+        };
+
+        assert_tokens(graph.codegen(), expected);
+    }
+}