1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
//! # Relu
//!
//! Applies the Rectified Linear Unit (ReLU) activation function element-wise.
//!
//! **ONNX Spec**: <https://onnx.ai/onnx/operators/onnx__Relu.html>
//!
//! ## Type Constraints
//! - `T`: Constrained to numeric tensors (float16, float32, float64, bfloat16, int8,
//! int16, int32, int64, uint8, uint16, uint32, uint64)
//!
//! ## Opset Versions
//! - **Opset 1-5**: Initial version
//! - **Opset 6-12**: Improved shape inference
//! - **Opset 13**: Expanded type support
//! - **Opset 14+**: Added bfloat16 support
use crate::ir::{Argument, Node, RawNode};
use crate::processor::{
InputSpec, NodeProcessor, NodeSpec, OutputPreferences, OutputSpec, ProcessError,
};
use onnx_ir_derive::NodeBuilder;
/// Node representation for Relu operation
#[derive(Debug, Clone, NodeBuilder)]
pub struct ReluNode {
pub name: String,
pub inputs: Vec<Argument>,
pub outputs: Vec<Argument>,
}
pub(crate) struct ReluProcessor;
impl NodeProcessor for ReluProcessor {
type Config = ();
fn spec(&self) -> NodeSpec {
NodeSpec {
min_opset: 1,
max_opset: None,
inputs: InputSpec::Exact(1),
outputs: OutputSpec::Exact(1),
}
}
fn infer_types(
&self,
node: &mut RawNode,
_opset: usize,
_output_preferences: &OutputPreferences,
) -> Result<(), ProcessError> {
// TODO: Missing test coverage for different data types
// Tests only use f32. Spec supports float16, float64, bfloat16, and integer types.
// Add tests: relu_float64, relu_int32, relu_int64
// Note: Integer types may require checking if Burn backend supports ReLU on integers.
// TODO: Missing test coverage for edge cases
// No tests for:
// - All negative inputs (output should be all zeros)
// - All positive inputs (output should equal input)
// - Mixed with exact zeros
// - Very large/small values (numerical stability)
// Add tests: relu_all_negative, relu_all_positive, relu_with_zeros
// TODO: Missing test coverage for different tensor ranks
// Tests cover 2D and 4D. Add coverage for 1D, 3D, 5D tensors.
// Add tests: relu_1d, relu_3d, relu_5d
// Validate no unexpected attributes
if !node.attrs.is_empty() {
let keys: Vec<String> = node.attrs.keys().cloned().collect();
return Err(ProcessError::InvalidAttribute {
name: keys[0].clone(),
reason: format!("Relu does not support any attributes, found: {:?}", keys),
});
}
// Output type is same as input
crate::processor::same_as_input(node);
Ok(())
}
fn build_node(&self, builder: RawNode, _opset: usize) -> Node {
Node::Relu(ReluNode {
name: builder.name,
inputs: builder.inputs,
outputs: builder.outputs,
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ir::{ArgType, NodeType};
use crate::node::test_utils::TestNodeBuilder;
use burn_tensor::DType;
fn create_test_node() -> RawNode {
TestNodeBuilder::new(NodeType::Relu, "test_relu")
.input_tensor_f32("X", 4, Some(vec![1, 3, 224, 224]))
.output_tensor_f32("Y", 0, None) // Rank will be inferred
.build()
}
#[test]
fn test_relu_type_inference() {
let mut node = create_test_node();
let processor = ReluProcessor;
let prefs = OutputPreferences::new();
processor.infer_types(&mut node, 16, &prefs).unwrap();
// Check output type matches input
match &node.outputs[0].ty {
ArgType::Tensor(tensor) => {
assert_eq!(tensor.dtype, DType::F32);
assert_eq!(tensor.rank, 4);
assert_eq!(
tensor.static_shape,
Some(vec![Some(1), Some(3), Some(224), Some(224)])
);
}
_ => panic!("Expected tensor output"),
}
}
#[test]
fn test_relu_no_attributes_allowed() {
let mut node = create_test_node();
node.attrs.insert(
"invalid_attr".to_string(),
crate::ir::AttributeValue::Float32(0.5),
);
let processor = ReluProcessor;
let prefs = OutputPreferences::new();
let result = processor.infer_types(&mut node, 16, &prefs);
assert!(matches!(result, Err(ProcessError::InvalidAttribute { .. })));
}
#[test]
fn test_relu_shape_preservation() {
let mut node = TestNodeBuilder::new(NodeType::Relu, "test_relu")
.input_tensor_f32("X", 2, Some(vec![10, 20]))
.output_tensor_f32("Y", 0, None)
.build();
let processor = ReluProcessor;
let prefs = OutputPreferences::new();
processor.infer_types(&mut node, 16, &prefs).unwrap();
if let ArgType::Tensor(tensor) = &node.outputs[0].ty {
assert_eq!(tensor.static_shape, Some(vec![Some(10), Some(20)]));
} else {
panic!("Expected tensor output");
}
}
}