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
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
//! # Random Tensor Generation (RandomNormal, RandomUniform)
//!
//! This module implements processors for ONNX random tensor generation operations that create
//! tensors with random values sampled from specified distributions.
//!
//! **ONNX Specs**:
//! - RandomNormal: <https://onnx.ai/onnx/operators/onnx__RandomNormal.html>
//! - RandomUniform: <https://onnx.ai/onnx/operators/onnx__RandomUniform.html>
//!
//! ## Supported Operations
//!
//! - **RandomNormal**: Generates a tensor with values drawn from a normal distribution
//! - **RandomUniform**: Generates a tensor with values drawn from a uniform distribution
//!
//! Both operations require a `shape` attribute and support an optional `dtype` attribute
//! (defaults to FLOAT if not specified).
//!
//! ## Opset Versions
//!
//! ### RandomNormal
//! - **Opset 1**: Initial version with shape, dtype, mean, scale, and seed attributes.
//!
//! ### RandomUniform
//! - **Opset 1**: Initial version with shape, dtype, high, low, and seed attributes.
use onnx_ir_derive::NodeBuilder;
use crate::ir::Argument;
use crate::ir::{ArgType, DType, Node, RawNode, TensorType};
use crate::processor::{
InputSpec, NodeProcessor, NodeSpec, OutputPreferences, OutputSpec, ProcessError,
};
use crate::protos::tensor_proto::DataType;
use protobuf::Enum;
/// Configuration for RandomNormal operation.
#[derive(Debug, Clone)]
pub struct RandomNormalConfig {
pub mean: f64,
pub scale: f64,
pub shape: Vec<usize>,
}
/// Configuration for RandomUniform operation.
#[derive(Debug, Clone)]
pub struct RandomUniformConfig {
pub low: f64,
pub high: f64,
pub shape: Vec<usize>,
}
/// Enum config that can hold either RandomNormal or RandomUniform config
#[derive(Debug, Clone)]
pub enum RandomConfig {
Normal(RandomNormalConfig),
Uniform(RandomUniformConfig),
}
/// Node representation for RandomNormal operation
#[derive(Debug, Clone, NodeBuilder)]
pub struct RandomNormalNode {
pub name: String,
pub inputs: Vec<Argument>,
pub outputs: Vec<Argument>,
pub config: RandomNormalConfig,
}
/// Node representation for RandomUniform operation
#[derive(Debug, Clone, NodeBuilder)]
pub struct RandomUniformNode {
pub name: String,
pub inputs: Vec<Argument>,
pub outputs: Vec<Argument>,
pub config: RandomUniformConfig,
}
pub(crate) struct RandomProcessor;
impl NodeProcessor for RandomProcessor {
type Config = RandomConfig;
fn spec(&self) -> NodeSpec {
NodeSpec {
min_opset: 1,
max_opset: None,
inputs: InputSpec::Exact(0),
outputs: OutputSpec::Exact(1),
}
}
fn infer_types(
&self,
node: &mut RawNode,
_opset: usize,
_output_preferences: &OutputPreferences,
) -> Result<(), ProcessError> {
// TODO: Validate that this node has zero inputs (Random operations don't take inputs)
// Random operators should have no inputs, but this isn't validated.
// Should add: validate_input_count(node, 0) or validate_max_inputs(node, 0)
// Location: After validate_output_count
// TODO: RandomNormal mean and scale attributes not validated or extracted
// ONNX spec defines:
// - mean (float, default=0.0): Mean of the normal distribution
// - scale (float, default=1.0): Standard deviation of the normal distribution
// - seed (float, optional): Random seed for reproducibility
// These attributes exist in spec but are completely ignored by implementation.
// Should extract into config and validate ranges (scale > 0).
// Location: extract_config method
// TODO: RandomUniform high/low attributes not validated or extracted
// ONNX spec defines:
// - high (float, default=1.0): Upper boundary of uniform distribution (exclusive)
// - low (float, default=0.0): Lower boundary of uniform distribution (inclusive)
// - seed (float, optional): Random seed for reproducibility
// These attributes are ignored. Should extract into config and validate low < high.
// Location: extract_config method
// TODO: Missing test coverage for non-default distribution parameters
// Tests only validate basic shape output, not distribution parameters (mean, scale, high, low).
// Add tests: random_normal_custom_mean_scale, random_uniform_custom_range
// Note: Testing actual distribution properties is hard, but should at least extract attrs.
// TODO: Missing validation for seed attribute
// Spec mentions seed attribute for reproducibility but it's not extracted or validated.
// Add tests: random_normal_with_seed, random_uniform_with_seed
let dtype = node
.attrs
.get("dtype")
.map(|val| DataType::from_i32(val.clone().into_i32()).unwrap())
.unwrap_or(DataType::FLOAT);
let shape = node
.attrs
.get("shape")
.ok_or_else(|| ProcessError::Custom("required shape attribute missing".to_string()))?
.clone()
.into_i64s();
let elem_type = match dtype {
DataType::FLOAT => DType::F32,
DataType::DOUBLE => DType::F64,
_ => {
return Err(ProcessError::InvalidAttribute {
name: "dtype".to_string(),
reason: format!("tensor with type {dtype:?} not supported for random output"),
});
}
};
let rank = shape.len();
node.outputs[0].ty = ArgType::Tensor(TensorType {
dtype: elem_type,
rank,
static_shape: None,
});
Ok(())
}
fn extract_config(&self, node: &RawNode, _opset: usize) -> Result<Self::Config, ProcessError> {
let shape = node
.attrs
.get("shape")
.ok_or_else(|| ProcessError::Custom("required shape attribute missing".to_string()))?
.clone()
.into_i64s();
let shape: Vec<usize> = shape.into_iter().map(|i| i as usize).collect();
let config = match node.node_type {
crate::ir::NodeType::RandomNormal => {
let mean = node
.attrs
.get("mean")
.map(|v| v.clone().into_f32() as f64)
.unwrap_or(0.0);
let scale = node
.attrs
.get("scale")
.map(|v| v.clone().into_f32() as f64)
.unwrap_or(1.0);
RandomConfig::Normal(RandomNormalConfig { mean, scale, shape })
}
crate::ir::NodeType::RandomUniform => {
let low = node
.attrs
.get("low")
.map(|v| v.clone().into_f32() as f64)
.unwrap_or(0.0);
let high = node
.attrs
.get("high")
.map(|v| v.clone().into_f32() as f64)
.unwrap_or(1.0);
RandomConfig::Uniform(RandomUniformConfig { low, high, shape })
}
_ => {
return Err(ProcessError::Custom(format!(
"RandomProcessor does not support node type {:?}",
node.node_type
)));
}
};
Ok(config)
}
fn build_node(&self, builder: RawNode, opset: usize) -> Node {
let config = self
.extract_config(&builder, opset)
.expect("Config extraction failed");
match config {
RandomConfig::Normal(normal_config) => Node::RandomNormal(RandomNormalNode {
name: builder.name,
inputs: builder.inputs,
outputs: builder.outputs,
config: normal_config,
}),
RandomConfig::Uniform(uniform_config) => Node::RandomUniform(RandomUniformNode {
name: builder.name,
inputs: builder.inputs,
outputs: builder.outputs,
config: uniform_config,
}),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ir::NodeType;
use crate::node::test_utils::TestNodeBuilder;
use crate::protos::tensor_proto::DataType;
fn create_test_node(dtype: i32, shape: Vec<i64>) -> RawNode {
TestNodeBuilder::new(NodeType::RandomNormal, "test_random")
.output_tensor_f32("output", 0, None) // Rank 0 will be updated
.attr_int("dtype", dtype as i64)
.attr_ints("shape", shape)
.build()
}
#[test]
fn test_random_normal_float() {
let mut node = create_test_node(DataType::FLOAT.value(), vec![2, 3, 4]);
let processor = RandomProcessor;
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::F32);
assert_eq!(tensor.rank, 3);
}
_ => panic!("Expected tensor output"),
}
}
#[test]
fn test_random_normal_double() {
let mut node = create_test_node(DataType::DOUBLE.value(), vec![5]);
let processor = RandomProcessor;
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::F64);
assert_eq!(tensor.rank, 1);
}
_ => panic!("Expected tensor output"),
}
}
#[test]
fn test_random_normal_missing_shape() {
// Create node and then manually remove the shape attribute
let mut node = create_test_node(DataType::FLOAT.value(), vec![2, 3]);
node.attrs.remove("shape");
let processor = RandomProcessor;
let prefs = OutputPreferences::new();
let result = processor.infer_types(&mut node, 16, &prefs);
assert!(matches!(result, Err(ProcessError::Custom(_))));
}
#[test]
fn test_random_normal_unsupported_type() {
let mut node = create_test_node(DataType::INT32.value(), vec![2, 3]);
let processor = RandomProcessor;
let prefs = OutputPreferences::new();
let result = processor.infer_types(&mut node, 16, &prefs);
assert!(matches!(result, Err(ProcessError::InvalidAttribute { .. })));
}
}