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
//! # NonZero
//!
//! Returns indices of non-zero elements.
//!
//! **ONNX Spec**: <https://onnx.ai/onnx/operators/onnx__NonZero.html>
//!
//! ## Opset Versions
//! - **Opset 9**: Initial version. Returns 2D tensor with shape [rank(X), num_non_zero].
//! - **Opset 13**: Added support for bfloat16 input type.
//!
//! ## Type Constraints (from ONNX spec)
//! - T: tensor(uint8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(int8), tensor(int16),
//! tensor(int32), tensor(int64), tensor(bfloat16), tensor(float16), tensor(float), tensor(double),
//! tensor(string), tensor(bool), tensor(complex64), tensor(complex128)
//!
//! TODO: Add validation for supported input types - current implementation accepts any tensor type
//! without validation against ONNX type constraints. While this may work for numeric types,
//! unsupported types like string, complex64, complex128 should be explicitly rejected.
//! Location: infer_types method after line 38
use onnx_ir_derive::NodeBuilder;
use crate::ir::{ArgType, Argument, DType, Node, RawNode, TensorType};
use crate::processor::{
InputSpec, NodeProcessor, NodeSpec, OutputPreferences, OutputSpec, ProcessError,
};
/// Node representation for NonZero operation
#[derive(Debug, Clone, NodeBuilder)]
pub struct NonZeroNode {
pub name: String,
pub inputs: Vec<Argument>,
pub outputs: Vec<Argument>,
}
pub(crate) struct NonZeroProcessor;
impl NodeProcessor for NonZeroProcessor {
type Config = ();
fn spec(&self) -> NodeSpec {
NodeSpec {
min_opset: 9,
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> {
// Note: Implementation correctly validates inputs/outputs per spec (1 input, 1 output)
match &node.inputs[0].ty {
ArgType::Tensor(_tensor) => {
// TODO: Missing test coverage for zero-size input tensors (e.g., shape [0, 5] or [3, 0])
// The spec allows zero-size tensors, but there's no test validating behavior when
// input tensor has zero elements. Should return [rank, 0] shaped output.
// Add test: nonzero_zero_size_tensor
// Output is always a 2D Int64 tensor
// Shape: [input_tensor_rank, num_nonzero_elements]
// First dimension equals input tensor rank
// Second dimension is dynamic (depends on data)
node.outputs[0].ty = ArgType::Tensor(TensorType {
dtype: DType::I64,
rank: 2,
static_shape: None, // Dynamic shape - second dimension depends on number of nonzero elements
});
}
_ => {
return Err(ProcessError::TypeMismatch {
expected: "Tensor".to_string(),
actual: format!("{:?}", node.inputs[0].ty),
});
}
}
Ok(())
}
fn build_node(&self, builder: RawNode, _opset: usize) -> Node {
Node::NonZero(NonZeroNode {
name: builder.name,
inputs: builder.inputs,
outputs: builder.outputs,
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ir::{DType, NodeType};
use crate::node::test_utils::TestNodeBuilder;
#[test]
fn test_nonzero_update_output() {
let mut node = TestNodeBuilder::new(NodeType::NonZero, "test_nonzero")
.input_tensor_f32("input", 3, Some(vec![2, 3, 4]))
.output_tensor_i64("output", 2, None) // rank will be updated
.build();
let processor = NonZeroProcessor;
let prefs = OutputPreferences::new();
processor.infer_types(&mut node, 16, &prefs).unwrap();
match &node.outputs[0].ty {
ArgType::Tensor(tensor) => {
assert_eq!(tensor.dtype, DType::I64);
assert_eq!(tensor.rank, 2);
assert_eq!(tensor.static_shape, None); // Dynamic shape
}
_ => panic!("Expected tensor output"),
}
}
#[test]
fn test_nonzero_update_output_1d() {
let mut node = TestNodeBuilder::new(NodeType::NonZero, "test_nonzero_1d")
.input_tensor_i32("input", 1, Some(vec![5]))
.output_tensor_i64("output", 2, None)
.build();
let processor = NonZeroProcessor;
let prefs = OutputPreferences::new();
processor.infer_types(&mut node, 16, &prefs).unwrap();
match &node.outputs[0].ty {
ArgType::Tensor(tensor) => {
assert_eq!(tensor.dtype, DType::I64);
assert_eq!(tensor.rank, 2);
assert_eq!(tensor.static_shape, None); // Dynamic shape
}
_ => panic!("Expected tensor output"),
}
}
#[test]
fn test_nonzero_update_output_4d() {
let mut node = TestNodeBuilder::new(NodeType::NonZero, "test_nonzero_4d")
.input_tensor_f64("input", 4, Some(vec![2, 3, 4, 5]))
.output_tensor_i64("output", 2, None)
.build();
let processor = NonZeroProcessor;
let prefs = OutputPreferences::new();
processor.infer_types(&mut node, 16, &prefs).unwrap();
match &node.outputs[0].ty {
ArgType::Tensor(tensor) => {
assert_eq!(tensor.dtype, DType::I64);
assert_eq!(tensor.rank, 2);
assert_eq!(tensor.static_shape, None); // Dynamic shape
}
_ => panic!("Expected tensor output"),
}
}
}