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
//! # IsNaN
//!
//! Returns which elements of the input are NaN (Not a Number).
//!
//! **ONNX Spec**: <https://onnx.ai/onnx/operators/onnx__IsNaN.html>
//!
//! ## Type Constraints
//!
//! - T1: tensor(float16), tensor(float), tensor(double), tensor(bfloat16), tensor(float8e4m3fn),
//! tensor(float8e4m3fnuz), tensor(float8e5m2), tensor(float8e5m2fnuz)
//! - T2: tensor(bool)
//!
//! ## Opset Versions
//! - **Opset 9-12**: Initial version
//! - **Opset 13-19**: Extended type support (added bfloat16)
//! - **Opset 20+**: Added float8 type variants support
//!
//! ## Missing Test Coverage
//! - TODO: No test for mixed NaN/Inf/finite values in same tensor - Current test only has NaN and finite
//! - TODO: No test for zero-size tensors - Edge case for empty tensor handling
//! - TODO: No test validating that input must be floating-point type - Integer inputs should be rejected
//! - TODO: No test for higher-rank tensors (3D, 4D) - Only 2D tensor tested
//! - TODO: No test for positive/negative NaN variants - Some platforms distinguish signaling/quiet NaN
use onnx_ir_derive::NodeBuilder;
use crate::ir::{Argument, Node, RawNode};
use crate::processor::{
InputSpec, NodeProcessor, NodeSpec, OutputPreferences, OutputSpec, ProcessError,
};
/// Node representation for IsNaN operation
#[derive(Debug, Clone, NodeBuilder)]
pub struct IsNaNNode {
pub name: String,
pub inputs: Vec<Argument>,
pub outputs: Vec<Argument>,
}
pub(crate) struct IsNaNProcessor;
impl NodeProcessor for IsNaNProcessor {
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> {
// TODO: Validate input dtype is floating-point - Type constraint T1: tensor(float16), tensor(float), tensor(double), tensor(bfloat16), tensor(float8*) not enforced - Integer inputs should be rejected - burn/crates/onnx-ir/src/node/is_nan.rs:42
// TODO: Validate that no unexpected attributes are present
// The spec states "None" for attributes
if let Some((key, _value)) = node.attrs.iter().next() {
return Err(ProcessError::InvalidAttribute {
name: key.clone(),
reason: format!("IsNaN does not accept any attributes, found: {}", key),
});
}
// Output is boolean tensor with same shape as input
crate::node::comparison::elementwise_comparison_outputs(node);
Ok(())
}
fn build_node(&self, builder: RawNode, _opset: usize) -> Node {
Node::IsNaN(IsNaNNode {
name: builder.name,
inputs: builder.inputs,
outputs: builder.outputs,
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ir::{ArgType, BoolStore, DType, NodeType};
use crate::node::test_utils::TestNodeBuilder;
#[test]
fn test_is_nan_basic() {
let mut node = TestNodeBuilder::new(NodeType::IsNaN, "test_is_nan")
.input_tensor_f32("data", 4, None)
.output_tensor_bool("output", 4, None)
.build();
let processor = IsNaNProcessor;
let prefs = OutputPreferences::new();
processor.infer_types(&mut node, 16, &prefs).unwrap();
// Output should be boolean with same rank as input
match &node.outputs[0].ty {
ArgType::Tensor(tensor) => {
assert_eq!(tensor.dtype, DType::Bool(BoolStore::Native));
assert_eq!(tensor.rank, 4);
}
_ => panic!("Expected tensor output"),
}
}
#[test]
fn test_is_nan_scalar() {
let mut node = TestNodeBuilder::new(NodeType::IsNaN, "test_is_nan")
.add_input("data", ArgType::ScalarNative(DType::F32))
.add_output(
"output",
ArgType::ScalarNative(DType::Bool(BoolStore::Native)),
)
.build();
let processor = IsNaNProcessor;
let prefs = OutputPreferences::new();
processor.infer_types(&mut node, 16, &prefs).unwrap();
// Output should be boolean scalar
match &node.outputs[0].ty {
ArgType::ScalarNative(elem_type) => {
assert_eq!(*elem_type, DType::Bool(BoolStore::Native));
}
_ => panic!("Expected scalar output"),
}
}
}