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
158
159
//! # InstanceNormalization
//!
//! Applies instance normalization to the input as described in
//! <https://arxiv.org/abs/1607.08022>.
//!
//! The operation normalizes each channel in each data instance independently:
//! `y = scale * (x - mean) / sqrt(variance + epsilon) + B`, where mean and
//! variance are computed per instance per channel.
//!
//! **ONNX Spec**: <https://onnx.ai/onnx/operators/onnx__InstanceNormalization.html>
//!
//! ## Type Constraints
//! - T: tensor(float16), tensor(float), tensor(double), tensor(bfloat16)
//!
//! ## Opset Versions
//! - **Opset 1-5**: Earlier versions with different epsilon handling
//! - **Opset 6+**: Current version with epsilon=1e-5 default and standardized behavior
//!
//! ## Missing Test Coverage
//! - TODO: No test for custom epsilon values (e.g., epsilon=1e-3) - Only default epsilon tested
//! - TODO: No test for edge cases: zero-mean inputs, constant inputs, single channel
//! - TODO: No test validating behavior with different batch sizes or spatial dimensions
use derive_new::new;
use onnx_ir_derive::NodeBuilder;
use crate::ir::{Argument, Node, RawNode};
use crate::processor::{
InputSpec, NodeProcessor, NodeSpec, OutputPreferences, OutputSpec, ProcessError,
};
/// Configuration for InstanceNorm operations
#[derive(Debug, Clone, new)]
pub struct InstanceNormConfig {
/// Small constant added for numerical stability
pub epsilon: f64,
}
/// Node representation for InstanceNormalization operation
#[derive(Debug, Clone, NodeBuilder)]
pub struct InstanceNormalizationNode {
pub name: String,
pub inputs: Vec<Argument>,
pub outputs: Vec<Argument>,
pub config: InstanceNormConfig,
}
pub(crate) struct InstanceNormProcessor;
impl NodeProcessor for InstanceNormProcessor {
type Config = InstanceNormConfig;
fn spec(&self) -> NodeSpec {
NodeSpec {
min_opset: 1,
max_opset: None,
inputs: InputSpec::Exact(3),
outputs: OutputSpec::Exact(1),
}
}
fn lift_constants(&self, node: &mut RawNode, _opset: usize) -> Result<(), ProcessError> {
// Lift scale (input 1) and bias (input 2)
if node.inputs.len() > 1 && node.inputs[1].is_constant() {
node.inputs[1].to_static()?;
}
if node.inputs.len() > 2 && node.inputs[2].is_constant() {
node.inputs[2].to_static()?;
}
Ok(())
}
fn infer_types(
&self,
node: &mut RawNode,
_opset: usize,
_output_preferences: &OutputPreferences,
) -> Result<(), ProcessError> {
// TODO: Validate input tensor dtype is floating-point type - Type constraint T: tensor(float16), tensor(float), tensor(double), tensor(bfloat16) not enforced - burn/crates/onnx-ir/src/node/instance_norm.rs:88
// TODO: Validate that scale and bias tensors are 1D and have size C matching the channel dimension of input - Shape mismatch could cause runtime errors - burn/crates/onnx-ir/src/node/instance_norm.rs:88
// TODO: Validate that input tensor is at least 3D (N x C x D1 ...) - Spec requires minimum rank of 3 - burn/crates/onnx-ir/src/node/instance_norm.rs:88
// Validate attributes before extracting config
for (key, _value) in node.attrs.iter() {
match key.as_str() {
"epsilon" => {}
_ => {
return Err(ProcessError::InvalidAttribute {
name: key.clone(),
reason: format!("Unexpected attribute for InstanceNorm: {key}"),
});
}
}
}
// Output type is same as input
crate::processor::same_as_input(node);
Ok(())
}
fn extract_config(&self, node: &RawNode, _opset: usize) -> Result<Self::Config, ProcessError> {
let mut epsilon = 1e-5;
for (key, value) in node.attrs.iter() {
if key.as_str() == "epsilon" {
// TODO: Validate epsilon > 0 for numerical stability
epsilon = value.clone().into_f32()
}
}
let config = InstanceNormConfig::new(epsilon as f64);
Ok(config)
}
fn build_node(&self, builder: RawNode, opset: usize) -> Node {
let config = self
.extract_config(&builder, opset)
.expect("Config extraction failed");
Node::InstanceNormalization(InstanceNormalizationNode {
name: builder.name,
inputs: builder.inputs,
outputs: builder.outputs,
config,
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ir::NodeType;
use crate::node::test_utils::TestNodeBuilder;
fn create_test_node(epsilon: f32, num_features: usize) -> TestNodeBuilder {
let weight_data = vec![1.0; num_features]; // Not important for the test
let bias_data = vec![0.0; num_features]; // Not important for the test
TestNodeBuilder::new(NodeType::InstanceNormalization, "test_instancenorm")
.input_tensor_f32("X", 3, None)
.input_tensor_f32_data("scale", weight_data, vec![num_features])
.input_tensor_f32_data("bias", bias_data, vec![num_features])
.output_tensor_f32("output", 3, None)
.attr_float("epsilon", epsilon)
}
#[test]
fn test_instance_norm_config_basic() {
let node = create_test_node(1e-5, 64).build_with_graph_data(16);
let mut node = node;
let processor = InstanceNormProcessor;
let prefs = OutputPreferences::new();
let config = processor.extract_config(&node, 16).unwrap();
processor.infer_types(&mut node, 16, &prefs).unwrap();
assert!(f64::abs(config.epsilon - 1e-5) < 1e-6);
}
}