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
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
//! # Softmax
//!
//! Applies the Softmax activation function along a specified axis.
//!
//! **ONNX Spec**: <https://onnx.ai/onnx/operators/onnx__Softmax.html>
//!
//! ## Type Constraints
//! - T: tensor(float16), tensor(float), tensor(double), tensor(bfloat16)
//!
//! ## Opset Versions
//! - **Opset 1**: Initial version with axis=1 default, operates on 2D tensors.
//! - **Opset 11**: Changed default axis to -1 (last dimension). Maintains backward compatibility with 2D coercion behavior.
//! - **Opset 13**: Removed 2D coercion behavior. Softmax now operates along specified axis directly without reshaping. This is the current behavior.
//!
//! **Implementation Note**: This implementation requires opset 13+ and uses the modern behavior (no 2D coercion). The axis attribute defaults to -1 as per opset 11+ specification.
use crate::ir::{ArgType, Argument, Node, RawNode};
use crate::processor::{
InputSpec, NodeProcessor, NodeSpec, OutputPreferences, OutputSpec, ProcessError,
};
use derive_new::new;
use onnx_ir_derive::NodeBuilder;
/// Configuration for Softmax operations
#[derive(Debug, Clone, new)]
pub struct SoftmaxConfig {
/// Axis along which to apply softmax
pub axis: usize,
}
/// Node representation for Softmax operation
#[derive(Debug, Clone, NodeBuilder)]
pub struct SoftmaxNode {
pub name: String,
pub inputs: Vec<Argument>,
pub outputs: Vec<Argument>,
pub config: SoftmaxConfig,
}
pub(crate) struct SoftmaxProcessor;
impl NodeProcessor for SoftmaxProcessor {
type Config = SoftmaxConfig;
fn spec(&self) -> NodeSpec {
NodeSpec {
min_opset: 13,
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> {
// FIXME: The spec requires the input rank to be >= 1 for the axis attribute to be valid.
// The implementation should validate that the tensor rank is at least 1.
// Edge case: what happens with a scalar (rank-0) input? Should be rejected.
// TODO: Missing validation that axis is in valid range [-rank, rank-1].
// Out-of-bounds axis values (after negative index handling) should be rejected.
// Infer output type
crate::processor::same_as_input(node);
Ok(())
}
fn extract_config(&self, node: &RawNode, _opset: usize) -> Result<Self::Config, ProcessError> {
// Extract the shape of the input tensor
let tensor = match &node.inputs.first().unwrap().ty {
ArgType::Tensor(tensor) => tensor.clone(),
_ => {
return Err(ProcessError::TypeMismatch {
expected: "Tensor".to_string(),
actual: format!("{:?}", node.inputs.first().unwrap().ty),
});
}
};
// Extract the axis attribute (default: -1 per ONNX spec)
let mut axis: i64 = -1;
for (key, value) in node.attrs.iter() {
if key.as_str() == "axis" {
axis = value.clone().into_i64()
}
}
// if axis is negative, it is counted from the end
if axis < 0 {
axis += tensor.rank as i64;
}
let config = SoftmaxConfig {
axis: axis as usize,
};
Ok(config)
}
fn build_node(&self, builder: RawNode, opset: usize) -> Node {
let config = self
.extract_config(&builder, opset)
.expect("Config extraction failed");
Node::Softmax(SoftmaxNode {
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(axis: i64, input_rank: usize) -> RawNode {
TestNodeBuilder::new(NodeType::Softmax, "test_softmax")
.input_tensor_f32("data", input_rank, None)
.output_tensor_f32("output", input_rank, None)
.attr_int("axis", axis)
.build()
}
#[test]
fn test_softmax_config_basic() {
let node = create_test_node(-1, 3);
let mut node = node;
let processor = SoftmaxProcessor;
let prefs = OutputPreferences::new();
let config = processor.extract_config(&node, 16).unwrap();
processor.infer_types(&mut node, 16, &prefs).unwrap();
assert_eq!(config.axis, 2); // -1 + 3 = 2 (last dimension)
}
#[test]
fn test_softmax_config_explicit_axis() {
let node = create_test_node(1, 3);
let mut node = node;
let processor = SoftmaxProcessor;
let prefs = OutputPreferences::new();
let config = processor.extract_config(&node, 16).unwrap();
processor.infer_types(&mut node, 16, &prefs).unwrap();
assert_eq!(config.axis, 1);
}
#[test]
fn test_softmax_config_multiple_inputs() {
let mut node = create_test_node(1, 3);
// Add an extra input
let extra_input = TestNodeBuilder::new(NodeType::Identity, "temp")
.input_tensor_f32("extra", 1, None)
.build()
.inputs
.pop()
.unwrap();
node.inputs.push(extra_input);
let processor = SoftmaxProcessor;
let spec = processor.spec();
let result = crate::processor::validate_node_spec(&node, 16, &spec);
assert!(matches!(
result,
Err(ProcessError::InvalidInputCount {
expected: 1,
actual: 2
})
));
}
// TODO: Missing test for scalar (rank-0) input - should be rejected as rank must be >= 1.
// TODO: Missing test for axis out of range - e.g., axis=5 for rank-3 tensor.
// TODO: Missing test for opset 13 behavior change - spec changed from 2D coercion to direct axis operation.
// Need test to verify opset < 13 is rejected and opset 13+ works correctly.
// TODO: Missing test for type constraints - Softmax only supports float types.
// Need test to verify integer input is rejected (or properly handled).
// TODO: Missing test for 1D tensor with axis=0 - simplest valid case not tested.
// TODO: Missing test for negative axis normalization - axis=-1 should work correctly.
// Current test has this but doesn't verify the actual behavior, only config extraction.
}