1use crate::Model;
10
11pub use onnx_runtime_shape_inference::{
12 DimExpr, InferenceContext, InferenceFn, InferenceRegistry, MergePolicy, NodeIo, ShapeData,
13 ShapeInferError as ShapeError, TypeInfo, TypedShape,
14};
15
16#[derive(Clone, Debug, Default, PartialEq, Eq)]
18pub struct ShapeInferenceResult {
19 pub inferred: usize,
21 pub unknown: usize,
23 pub warnings: Vec<String>,
26}
27
28pub fn infer_shapes(model: &mut Model) -> Result<ShapeInferenceResult, ShapeError> {
33 let registry = InferenceRegistry::default_registry();
34 infer_shapes_with_registry(model, ®istry)
35}
36
37pub fn infer_shapes_with_registry(
42 model: &mut Model,
43 registry: &InferenceRegistry,
44) -> Result<ShapeInferenceResult, ShapeError> {
45 let opsets = model.graph.opset_imports.clone();
46 let report = registry.infer_graph(&mut model.graph, &opsets, MergePolicy::Permissive)?;
47 Ok(ShapeInferenceResult {
48 inferred: report.num_resolved(),
49 unknown: report.num_unresolved(),
50 warnings: Vec::new(),
51 })
52}
53
54pub fn register_shape_inference(
60 registry: &mut InferenceRegistry,
61 domain: &str,
62 op_type: &str,
63 min_opset: u64,
64 handler: InferenceFn,
65) {
66 registry.register(domain, op_type, min_opset, handler);
67}
68
69#[cfg(test)]
70mod tests {
71 use super::*;
72 use crate::ir::{
73 Attribute, DataType, Dim, Graph, Node, NodeId, Shape, TensorData, ValueId, WeightRef,
74 };
75
76 fn i64_initializer(graph: &mut Graph, name: &str, values: &[i64]) -> ValueId {
77 let value =
78 graph.create_named_value(name, DataType::Int64, vec![Dim::Static(values.len())]);
79 let mut bytes = Vec::with_capacity(values.len().saturating_mul(size_of::<i64>()));
80 for item in values {
81 bytes.extend_from_slice(&item.to_le_bytes());
82 }
83 graph.set_initializer(
84 value,
85 WeightRef::Inline(TensorData::from_raw(
86 DataType::Int64,
87 vec![values.len()],
88 bytes,
89 )),
90 );
91 value
92 }
93
94 fn i64_scalar_initializer(graph: &mut Graph, name: &str, value: i64) -> ValueId {
95 let id = graph.create_named_value(name, DataType::Int64, Vec::<Dim>::new());
96 graph.set_initializer(
97 id,
98 WeightRef::Inline(TensorData::from_raw(
99 DataType::Int64,
100 Vec::new(),
101 value.to_le_bytes().to_vec(),
102 )),
103 );
104 id
105 }
106
107 #[test]
108 fn infers_matmul_and_add_output_shapes() {
109 let mut graph = Graph::new();
110 graph.opset_imports.insert(String::new(), 21);
111
112 let x =
113 graph.create_named_value("x", DataType::Float32, vec![Dim::Static(2), Dim::Static(3)]);
114 let weights = graph.create_named_value(
115 "weights",
116 DataType::Float32,
117 vec![Dim::Static(3), Dim::Static(4)],
118 );
119 let bias = graph.create_named_value("bias", DataType::Float32, vec![Dim::Static(4)]);
120 let product = graph.create_named_value("product", DataType::Float32, Shape::new());
121 let output = graph.create_named_value("output", DataType::Float32, Shape::new());
122
123 graph.add_input(x);
124 graph.add_input(weights);
125 graph.add_input(bias);
126 graph.insert_node(Node::new(
127 NodeId(0),
128 "MatMul",
129 vec![Some(x), Some(weights)],
130 vec![product],
131 ));
132 graph.insert_node(Node::new(
133 NodeId(1),
134 "Add",
135 vec![Some(product), Some(bias)],
136 vec![output],
137 ));
138 graph.add_output(output);
139
140 let mut model = Model::new(graph);
141 let result = infer_shapes(&mut model).unwrap();
142
143 assert_eq!(
144 model.graph.value(product).shape,
145 vec![Dim::Static(2), Dim::Static(4)]
146 );
147 assert_eq!(
148 model.graph.value(output).shape,
149 vec![Dim::Static(2), Dim::Static(4)]
150 );
151 assert_eq!(result.inferred, 5);
152 assert_eq!(result.unknown, 0);
153 assert!(result.warnings.is_empty());
154 }
155
156 fn infer_custom_identity(ctx: &mut InferenceContext) -> Result<(), ShapeError> {
157 if let Some(input) = ctx.input_type(0).cloned() {
158 ctx.set_output_type(0, input);
159 }
160 Ok(())
161 }
162
163 #[test]
164 fn custom_rule_uses_underlying_registry() {
165 let mut graph = Graph::new();
166 graph.opset_imports.insert("example".to_string(), 1);
167 let input = graph.create_named_value("input", DataType::Float32, vec![Dim::Static(7)]);
168 let output = graph.create_named_value("output", DataType::Float32, Shape::new());
169 graph.add_input(input);
170 let mut node = Node::new(NodeId(0), "CopyShape", vec![Some(input)], vec![output]);
171 node.domain = "example".to_string();
172 graph.insert_node(node);
173 graph.add_output(output);
174
175 let mut registry = InferenceRegistry::default_registry();
176 register_shape_inference(
177 &mut registry,
178 "example",
179 "CopyShape",
180 1,
181 infer_custom_identity,
182 );
183
184 let mut model = Model::new(graph);
185 let result = infer_shapes_with_registry(&mut model, ®istry).unwrap();
186
187 assert_eq!(model.graph.value(output).shape, vec![Dim::Static(7)]);
188 assert_eq!(result.inferred, 2);
189 assert_eq!(result.unknown, 0);
190 }
191
192 #[test]
193 fn infers_added_elementwise_expand_where_and_reduction_schemas() {
194 let mut graph = Graph::new();
195 graph.opset_imports.insert(String::new(), 18);
196
197 let x = graph.create_named_value(
198 "x",
199 DataType::Float32,
200 vec![Dim::Static(2), Dim::Static(3), Dim::Static(4)],
201 );
202 let exponent = graph.create_named_value("exponent", DataType::Int64, vec![Dim::Static(1)]);
203 let condition = graph.create_named_value(
204 "condition",
205 DataType::Bool,
206 vec![Dim::Static(2), Dim::Static(1), Dim::Static(4)],
207 );
208 let alternative = graph.create_named_value(
209 "alternative",
210 DataType::Float32,
211 vec![Dim::Static(1), Dim::Static(3), Dim::Static(1)],
212 );
213 let expand_source = graph.create_named_value(
214 "expand_source",
215 DataType::Float32,
216 vec![Dim::Static(3), Dim::Static(1)],
217 );
218 for value in [x, exponent, condition, alternative, expand_source] {
219 graph.add_input(value);
220 }
221
222 let expand_shape = i64_initializer(&mut graph, "expand_shape", &[2, 3, 4]);
223 let reduce_axis_1 = i64_initializer(&mut graph, "reduce_axis_1", &[1]);
224 let reduce_axis_2 = i64_initializer(&mut graph, "reduce_axis_2", &[2]);
225
226 let mut current = x;
227 for (id, op) in ["Sigmoid", "Tanh", "Erf", "Sqrt", "Exp", "Log", "Clip"]
228 .into_iter()
229 .enumerate()
230 {
231 let output =
232 graph.create_named_value(format!("{op}_out"), DataType::Float32, Shape::new());
233 let inputs = if op == "Clip" {
234 vec![Some(current), None, None]
235 } else {
236 vec![Some(current)]
237 };
238 graph.insert_node(Node::new(NodeId(id as u32), op, inputs, vec![output]));
239 current = output;
240 }
241
242 let pow = graph.create_named_value("pow", DataType::Float32, Shape::new());
243 graph.insert_node(Node::new(
244 NodeId(7),
245 "Pow",
246 vec![Some(current), Some(exponent)],
247 vec![pow],
248 ));
249 let selected = graph.create_named_value("selected", DataType::Float32, Shape::new());
250 graph.insert_node(Node::new(
251 NodeId(8),
252 "Where",
253 vec![Some(condition), Some(pow), Some(alternative)],
254 vec![selected],
255 ));
256 let expanded = graph.create_named_value("expanded", DataType::Float32, Shape::new());
257 graph.insert_node(Node::new(
258 NodeId(9),
259 "Expand",
260 vec![Some(expand_source), Some(expand_shape)],
261 vec![expanded],
262 ));
263 let sum = graph.create_named_value("sum", DataType::Float32, Shape::new());
264 let mut reduce_sum = Node::new(
265 NodeId(10),
266 "ReduceSum",
267 vec![Some(selected), Some(reduce_axis_1)],
268 vec![sum],
269 );
270 reduce_sum
271 .attributes
272 .insert("keepdims".into(), Attribute::Int(0));
273 graph.insert_node(reduce_sum);
274 let mean = graph.create_named_value("mean", DataType::Float32, Shape::new());
275 graph.insert_node(Node::new(
276 NodeId(11),
277 "ReduceMean",
278 vec![Some(selected), Some(reduce_axis_2)],
279 vec![mean],
280 ));
281 graph.add_output(expanded);
282 graph.add_output(sum);
283 graph.add_output(mean);
284
285 let mut model = Model::new(graph);
286 let result = infer_shapes(&mut model).unwrap();
287 assert_eq!(
288 model.graph.value(selected).shape,
289 vec![Dim::Static(2), Dim::Static(3), Dim::Static(4)]
290 );
291 assert_eq!(
292 model.graph.value(expanded).shape,
293 vec![Dim::Static(2), Dim::Static(3), Dim::Static(4)]
294 );
295 assert_eq!(
296 model.graph.value(sum).shape,
297 vec![Dim::Static(2), Dim::Static(4)]
298 );
299 assert_eq!(
300 model.graph.value(mean).shape,
301 vec![Dim::Static(2), Dim::Static(3), Dim::Static(1)]
302 );
303 assert_eq!(result.unknown, 0);
304 }
305
306 #[test]
307 fn infers_existing_conv_norm_and_movement_schemas() {
308 let mut graph = Graph::new();
309 graph.opset_imports.insert(String::new(), 21);
310
311 let image = graph.create_named_value(
312 "image",
313 DataType::Float32,
314 vec![
315 Dim::Static(1),
316 Dim::Static(3),
317 Dim::Static(32),
318 Dim::Static(32),
319 ],
320 );
321 let weights = graph.create_named_value(
322 "weights",
323 DataType::Float32,
324 vec![
325 Dim::Static(8),
326 Dim::Static(3),
327 Dim::Static(3),
328 Dim::Static(3),
329 ],
330 );
331 let scale = graph.create_named_value("scale", DataType::Float32, vec![Dim::Static(15_360)]);
332 let gather_indices = graph.create_named_value(
333 "gather_indices",
334 DataType::Int64,
335 vec![Dim::Static(5), Dim::Static(6)],
336 );
337 let concat_rhs = graph.create_named_value(
338 "concat_rhs",
339 DataType::Float32,
340 vec![
341 Dim::Static(1),
342 Dim::Static(32),
343 Dim::Static(7),
344 Dim::Static(6),
345 Dim::Static(32),
346 ],
347 );
348 for value in [image, weights, scale, gather_indices, concat_rhs] {
349 graph.add_input(value);
350 }
351
352 let starts = i64_initializer(&mut graph, "starts", &[0]);
353 let ends = i64_initializer(&mut graph, "ends", &[10]);
354 let axes = i64_initializer(&mut graph, "axes", &[2]);
355 let steps = i64_initializer(&mut graph, "steps", &[2]);
356 let reshape_target = i64_initializer(&mut graph, "reshape_target", &[2, -1]);
357
358 let convolved = graph.create_named_value("convolved", DataType::Float32, Shape::new());
359 let mut conv = Node::new(
360 NodeId(0),
361 "Conv",
362 vec![Some(image), Some(weights)],
363 vec![convolved],
364 );
365 conv.attributes
366 .insert("strides".into(), Attribute::Ints(vec![2, 2]));
367 conv.attributes
368 .insert("pads".into(), Attribute::Ints(vec![1, 1, 1, 1]));
369 graph.insert_node(conv);
370
371 let gathered = graph.create_named_value("gathered", DataType::Float32, Shape::new());
372 let mut gather = Node::new(
373 NodeId(1),
374 "Gather",
375 vec![Some(image), Some(gather_indices)],
376 vec![gathered],
377 );
378 gather.attributes.insert("axis".into(), Attribute::Int(1));
379 graph.insert_node(gather);
380
381 let transposed = graph.create_named_value("transposed", DataType::Float32, Shape::new());
382 let mut transpose = Node::new(
383 NodeId(2),
384 "Transpose",
385 vec![Some(gathered)],
386 vec![transposed],
387 );
388 transpose
389 .attributes
390 .insert("perm".into(), Attribute::Ints(vec![0, 3, 1, 2, 4]));
391 graph.insert_node(transpose);
392
393 let concatenated =
394 graph.create_named_value("concatenated", DataType::Float32, Shape::new());
395 let mut concat = Node::new(
396 NodeId(3),
397 "Concat",
398 vec![Some(transposed), Some(concat_rhs)],
399 vec![concatenated],
400 );
401 concat.attributes.insert("axis".into(), Attribute::Int(2));
402 graph.insert_node(concat);
403
404 let sliced = graph.create_named_value("sliced", DataType::Float32, Shape::new());
405 graph.insert_node(Node::new(
406 NodeId(4),
407 "Slice",
408 vec![
409 Some(concatenated),
410 Some(starts),
411 Some(ends),
412 Some(axes),
413 Some(steps),
414 ],
415 vec![sliced],
416 ));
417
418 let reshaped = graph.create_named_value("reshaped", DataType::Float32, Shape::new());
419 graph.insert_node(Node::new(
420 NodeId(5),
421 "Reshape",
422 vec![Some(sliced), Some(reshape_target)],
423 vec![reshaped],
424 ));
425 let normalized = graph.create_named_value("normalized", DataType::Float32, Shape::new());
426 let mean = graph.create_named_value("norm_mean", DataType::Float32, Shape::new());
427 let inv_std = graph.create_named_value("inv_std", DataType::Float32, Shape::new());
428 graph.insert_node(Node::new(
429 NodeId(6),
430 "LayerNormalization",
431 vec![Some(reshaped), Some(scale), None],
432 vec![normalized, mean, inv_std],
433 ));
434 let identity = graph.create_named_value("identity", DataType::Float32, Shape::new());
435 graph.insert_node(Node::new(
436 NodeId(7),
437 "Identity",
438 vec![Some(normalized)],
439 vec![identity],
440 ));
441 graph.add_output(convolved);
442 graph.add_output(identity);
443
444 let mut model = Model::new(graph);
445 infer_shapes(&mut model).unwrap();
446 assert_eq!(
447 model.graph.value(convolved).shape,
448 vec![
449 Dim::Static(1),
450 Dim::Static(8),
451 Dim::Static(16),
452 Dim::Static(16)
453 ]
454 );
455 assert_eq!(
456 model.graph.value(transposed).shape,
457 vec![
458 Dim::Static(1),
459 Dim::Static(32),
460 Dim::Static(5),
461 Dim::Static(6),
462 Dim::Static(32)
463 ]
464 );
465 assert_eq!(
466 model.graph.value(identity).shape,
467 vec![Dim::Static(2), Dim::Static(15_360)]
468 );
469 assert_eq!(
470 model.graph.value(mean).shape,
471 vec![Dim::Static(2), Dim::Static(1)]
472 );
473 }
474
475 #[test]
476 fn infers_round_four_normalization_reduction_and_arg_shapes() {
477 let mut graph = Graph::new();
478 graph.opset_imports.insert(String::new(), 24);
479 let x = graph.create_named_value(
480 "x",
481 DataType::Float32,
482 vec![Dim::Static(2), Dim::Static(3), Dim::Static(4)],
483 );
484 let scale = graph.create_named_value("scale", DataType::Float16, vec![Dim::Static(4)]);
485 graph.add_input(x);
486 graph.add_input(scale);
487 let axes = i64_initializer(&mut graph, "axes", &[1]);
488
489 let log_softmax = graph.create_named_value("log_softmax", DataType::Float32, Shape::new());
490 graph.insert_node(Node::new(
491 NodeId(0),
492 "LogSoftmax",
493 vec![Some(x)],
494 vec![log_softmax],
495 ));
496
497 let rms = graph.create_named_value("rms", DataType::Undefined, Shape::new());
498 graph.insert_node(Node::new(
499 NodeId(1),
500 "RMSNormalization",
501 vec![Some(x), Some(scale)],
502 vec![rms],
503 ));
504
505 let mut reduction_outputs = Vec::new();
506 for (offset, op) in [
507 "ReduceMax",
508 "ReduceMin",
509 "ReduceProd",
510 "ReduceL1",
511 "ReduceL2",
512 "ReduceLogSum",
513 "ReduceLogSumExp",
514 "ReduceSumSquare",
515 ]
516 .into_iter()
517 .enumerate()
518 {
519 let output =
520 graph.create_named_value(format!("{op}_out"), DataType::Float32, Shape::new());
521 let mut node = Node::new(
522 NodeId(2 + offset as u32),
523 op,
524 vec![Some(x), Some(axes)],
525 vec![output],
526 );
527 node.attributes.insert("keepdims".into(), Attribute::Int(0));
528 graph.insert_node(node);
529 reduction_outputs.push(output);
530 }
531
532 let arg_max = graph.create_named_value("arg_max", DataType::Int64, Shape::new());
533 let mut arg_max_node = Node::new(NodeId(10), "ArgMax", vec![Some(x)], vec![arg_max]);
534 arg_max_node
535 .attributes
536 .insert("axis".into(), Attribute::Int(-1));
537 arg_max_node
538 .attributes
539 .insert("keepdims".into(), Attribute::Int(0));
540 graph.insert_node(arg_max_node);
541
542 let arg_min = graph.create_named_value("arg_min", DataType::Int64, Shape::new());
543 let mut arg_min_node = Node::new(NodeId(11), "ArgMin", vec![Some(x)], vec![arg_min]);
544 arg_min_node
545 .attributes
546 .insert("axis".into(), Attribute::Int(1));
547 graph.insert_node(arg_min_node);
548
549 for output in [log_softmax, rms, arg_max, arg_min]
550 .into_iter()
551 .chain(reduction_outputs.iter().copied())
552 {
553 graph.add_output(output);
554 }
555
556 let mut model = Model::new(graph);
557 let result = infer_shapes(&mut model).unwrap();
558 assert_eq!(
559 model.graph.value(log_softmax).shape,
560 vec![Dim::Static(2), Dim::Static(3), Dim::Static(4)]
561 );
562 assert_eq!(
563 (
564 model.graph.value(rms).dtype,
565 model.graph.value(rms).shape.clone()
566 ),
567 (
568 DataType::Float16,
569 vec![Dim::Static(2), Dim::Static(3), Dim::Static(4)]
570 )
571 );
572 for output in reduction_outputs {
573 assert_eq!(
574 model.graph.value(output).shape,
575 vec![Dim::Static(2), Dim::Static(4)]
576 );
577 }
578 assert_eq!(
579 (
580 model.graph.value(arg_max).dtype,
581 model.graph.value(arg_max).shape.clone()
582 ),
583 (DataType::Int64, vec![Dim::Static(2), Dim::Static(3)])
584 );
585 assert_eq!(
586 (
587 model.graph.value(arg_min).dtype,
588 model.graph.value(arg_min).shape.clone()
589 ),
590 (
591 DataType::Int64,
592 vec![Dim::Static(2), Dim::Static(1), Dim::Static(4)]
593 )
594 );
595 assert_eq!(result.unknown, 0);
596 }
597
598 #[test]
599 fn infers_round_five_index_predicate_and_shape_ops() {
600 let mut graph = Graph::new();
601 graph.opset_imports.insert(String::new(), 25);
602
603 let data = graph.create_named_value(
604 "data",
605 DataType::Float32,
606 vec![Dim::Static(2), Dim::Static(3), Dim::Static(4)],
607 );
608 let gather_elements_indices = graph.create_named_value(
609 "gather_elements_indices",
610 DataType::Int64,
611 vec![Dim::Static(2), Dim::Static(5), Dim::Static(4)],
612 );
613 let gather_nd_indices = graph.create_named_value(
614 "gather_nd_indices",
615 DataType::Int64,
616 vec![Dim::Static(2), Dim::Static(5), Dim::Static(2)],
617 );
618 let lhs = graph.create_named_value(
619 "lhs",
620 DataType::Float32,
621 vec![Dim::Static(2), Dim::Static(1), Dim::Static(4)],
622 );
623 let rhs = graph.create_named_value(
624 "rhs",
625 DataType::Float32,
626 vec![Dim::Static(1), Dim::Static(3), Dim::Static(1)],
627 );
628 let split_input = graph.create_named_value(
629 "split_input",
630 DataType::Float32,
631 vec![Dim::Static(2), Dim::Static(6), Dim::Static(4)],
632 );
633 for input in [
634 data,
635 gather_elements_indices,
636 gather_nd_indices,
637 lhs,
638 rhs,
639 split_input,
640 ] {
641 graph.add_input(input);
642 }
643
644 let range_start = i64_scalar_initializer(&mut graph, "range_start", 0);
645 let range_limit = i64_scalar_initializer(&mut graph, "range_limit", 10);
646 let range_delta = i64_scalar_initializer(&mut graph, "range_delta", 2);
647 let split_sizes = i64_initializer(&mut graph, "split_sizes", &[2, 4]);
648
649 let gathered_elements =
650 graph.create_named_value("gathered_elements", DataType::Undefined, Shape::new());
651 let mut gather_elements = Node::new(
652 NodeId(0),
653 "GatherElements",
654 vec![Some(data), Some(gather_elements_indices)],
655 vec![gathered_elements],
656 );
657 gather_elements
658 .attributes
659 .insert("axis".into(), Attribute::Int(-2));
660 graph.insert_node(gather_elements);
661
662 let gathered_nd =
663 graph.create_named_value("gathered_nd", DataType::Undefined, Shape::new());
664 let mut gather_nd = Node::new(
665 NodeId(1),
666 "GatherND",
667 vec![Some(data), Some(gather_nd_indices)],
668 vec![gathered_nd],
669 );
670 gather_nd
671 .attributes
672 .insert("batch_dims".into(), Attribute::Int(1));
673 graph.insert_node(gather_nd);
674
675 let equal = graph.create_named_value("equal", DataType::Undefined, Shape::new());
676 graph.insert_node(Node::new(
677 NodeId(2),
678 "Equal",
679 vec![Some(lhs), Some(rhs)],
680 vec![equal],
681 ));
682 let greater = graph.create_named_value("greater", DataType::Undefined, Shape::new());
683 graph.insert_node(Node::new(
684 NodeId(3),
685 "Greater",
686 vec![Some(lhs), Some(rhs)],
687 vec![greater],
688 ));
689 let less = graph.create_named_value("less", DataType::Undefined, Shape::new());
690 graph.insert_node(Node::new(
691 NodeId(4),
692 "Less",
693 vec![Some(lhs), Some(rhs)],
694 vec![less],
695 ));
696 let and = graph.create_named_value("and", DataType::Undefined, Shape::new());
697 graph.insert_node(Node::new(
698 NodeId(5),
699 "And",
700 vec![Some(equal), Some(greater)],
701 vec![and],
702 ));
703 let or = graph.create_named_value("or", DataType::Undefined, Shape::new());
704 graph.insert_node(Node::new(
705 NodeId(6),
706 "Or",
707 vec![Some(and), Some(less)],
708 vec![or],
709 ));
710 let not = graph.create_named_value("not", DataType::Undefined, Shape::new());
711 graph.insert_node(Node::new(NodeId(7), "Not", vec![Some(or)], vec![not]));
712
713 let cast = graph.create_named_value("cast", DataType::Undefined, Shape::new());
714 let mut cast_node = Node::new(NodeId(8), "Cast", vec![Some(data)], vec![cast]);
715 cast_node.attributes.insert(
716 "to".into(),
717 Attribute::Int(i64::from(DataType::Int64.to_onnx())),
718 );
719 graph.insert_node(cast_node);
720
721 let shape = graph.create_named_value("shape", DataType::Undefined, Shape::new());
722 let mut shape_node = Node::new(NodeId(9), "Shape", vec![Some(data)], vec![shape]);
723 shape_node
724 .attributes
725 .insert("start".into(), Attribute::Int(1));
726 shape_node
727 .attributes
728 .insert("end".into(), Attribute::Int(i64::MAX));
729 graph.insert_node(shape_node);
730
731 let size = graph.create_named_value("size", DataType::Undefined, Shape::new());
732 graph.insert_node(Node::new(NodeId(10), "Size", vec![Some(data)], vec![size]));
733 let non_zero = graph.create_named_value("non_zero", DataType::Undefined, Shape::new());
734 graph.insert_node(Node::new(
735 NodeId(11),
736 "NonZero",
737 vec![Some(data)],
738 vec![non_zero],
739 ));
740 let range = graph.create_named_value("range", DataType::Undefined, Shape::new());
741 graph.insert_node(Node::new(
742 NodeId(12),
743 "Range",
744 vec![Some(range_start), Some(range_limit), Some(range_delta)],
745 vec![range],
746 ));
747
748 let split_first =
749 graph.create_named_value("split_first", DataType::Undefined, Shape::new());
750 let split_second =
751 graph.create_named_value("split_second", DataType::Undefined, Shape::new());
752 let mut split = Node::new(
753 NodeId(13),
754 "Split",
755 vec![Some(split_input), Some(split_sizes)],
756 vec![split_first, split_second],
757 );
758 split.attributes.insert("axis".into(), Attribute::Int(-2));
759 graph.insert_node(split);
760
761 for output in [
762 gathered_elements,
763 gathered_nd,
764 equal,
765 greater,
766 less,
767 and,
768 or,
769 not,
770 cast,
771 shape,
772 size,
773 non_zero,
774 range,
775 split_first,
776 split_second,
777 ] {
778 graph.add_output(output);
779 }
780
781 let mut model = Model::new(graph);
782 let result = infer_shapes(&mut model).unwrap();
783
784 assert_eq!(
785 (
786 model.graph.value(gathered_elements).dtype,
787 model.graph.value(gathered_elements).shape.clone()
788 ),
789 (
790 DataType::Float32,
791 vec![Dim::Static(2), Dim::Static(5), Dim::Static(4)]
792 )
793 );
794 assert_eq!(
795 model.graph.value(gathered_nd).shape,
796 vec![Dim::Static(2), Dim::Static(5)]
797 );
798 for output in [equal, greater, less, and, or, not] {
799 assert_eq!(
800 (
801 model.graph.value(output).dtype,
802 model.graph.value(output).shape.clone()
803 ),
804 (
805 DataType::Bool,
806 vec![Dim::Static(2), Dim::Static(3), Dim::Static(4)]
807 )
808 );
809 }
810 assert_eq!(
811 (
812 model.graph.value(cast).dtype,
813 model.graph.value(cast).shape.clone()
814 ),
815 (
816 DataType::Int64,
817 vec![Dim::Static(2), Dim::Static(3), Dim::Static(4)]
818 )
819 );
820 assert_eq!(
821 (
822 model.graph.value(shape).dtype,
823 model.graph.value(shape).shape.clone()
824 ),
825 (DataType::Int64, vec![Dim::Static(2)])
826 );
827 assert_eq!(
828 (
829 model.graph.value(size).dtype,
830 model.graph.value(size).shape.clone()
831 ),
832 (DataType::Int64, Vec::<Dim>::new())
833 );
834 assert!(matches!(
835 model.graph.value(non_zero).shape.as_slice(),
836 [Dim::Static(3), Dim::Symbolic(_)]
837 ));
838 assert_eq!(model.graph.value(range).shape, vec![Dim::Static(5)]);
839 assert_eq!(
840 model.graph.value(split_first).shape,
841 vec![Dim::Static(2), Dim::Static(2), Dim::Static(4)]
842 );
843 assert_eq!(
844 model.graph.value(split_second).shape,
845 vec![Dim::Static(2), Dim::Static(4), Dim::Static(4)]
846 );
847 assert_eq!(result.unknown, 0);
848 }
849
850 #[test]
851 fn round_five_shape_arithmetic_rejects_unrepresentable_values() {
852 let mut range_graph = Graph::new();
853 range_graph.opset_imports.insert(String::new(), 25);
854 let start = i64_scalar_initializer(&mut range_graph, "start", i64::MIN);
855 let limit = i64_scalar_initializer(&mut range_graph, "limit", i64::MAX);
856 let delta = i64_scalar_initializer(&mut range_graph, "delta", 1);
857 let output = range_graph.create_named_value("range", DataType::Undefined, Shape::new());
858 range_graph.insert_node(Node::new(
859 NodeId(0),
860 "Range",
861 vec![Some(start), Some(limit), Some(delta)],
862 vec![output],
863 ));
864 range_graph.add_output(output);
865 let error = infer_shapes(&mut Model::new(range_graph)).unwrap_err();
866 assert!(error.to_string().contains("exceeds isize::MAX"));
867
868 let mut split_graph = Graph::new();
869 split_graph.opset_imports.insert(String::new(), 25);
870 let input = split_graph.create_named_value(
871 "input",
872 DataType::Float32,
873 vec![Dim::Static(2), Dim::Static(4)],
874 );
875 let first = split_graph.create_named_value("first", DataType::Undefined, Shape::new());
876 let second = split_graph.create_named_value("second", DataType::Undefined, Shape::new());
877 split_graph.add_input(input);
878 let mut split = Node::new(NodeId(0), "Split", vec![Some(input)], vec![first, second]);
879 split
880 .attributes
881 .insert("axis".into(), Attribute::Int(i64::MIN));
882 split_graph.insert_node(split);
883 split_graph.add_output(first);
884 split_graph.add_output(second);
885 let error = infer_shapes(&mut Model::new(split_graph)).unwrap_err();
886 assert!(error.to_string().contains("axis"));
887 }
888
889 #[test]
890 fn if_inference_unions_branch_output_shapes() {
891 fn branch(shape: Vec<Dim>) -> Graph {
892 let mut graph = Graph::new();
893 let output = graph.create_named_value("branch_output", DataType::Float32, shape);
894 graph.add_input(output);
895 graph.add_output(output);
896 graph
897 }
898
899 let mut graph = Graph::new();
900 graph.opset_imports.insert(String::new(), 21);
901 let condition = graph.create_named_value("condition", DataType::Bool, Vec::<Dim>::new());
902 graph.add_input(condition);
903 let output = graph.create_named_value("output", DataType::Float32, Shape::new());
904 let node = graph.insert_node(Node::new(
905 NodeId(0),
906 "If",
907 vec![Some(condition)],
908 vec![output],
909 ));
910 graph.subgraphs.insert(
911 (node, "then_branch".into()),
912 branch(vec![Dim::Static(2), Dim::Static(4)]),
913 );
914 graph.subgraphs.insert(
915 (node, "else_branch".into()),
916 branch(vec![Dim::Static(3), Dim::Static(4)]),
917 );
918 graph.add_output(output);
919
920 let mut model = Model::new(graph);
921 let result = infer_shapes(&mut model).unwrap();
922 assert!(matches!(
923 model.graph.value(output).shape.as_slice(),
924 [Dim::Symbolic(_), Dim::Static(4)]
925 ));
926 assert_eq!(result.unknown, 0);
927 }
928
929 #[test]
930 fn if_inference_does_not_invent_unresolved_branch_shapes() {
931 fn unresolved_branch() -> Graph {
932 let mut graph = Graph::new();
933 let input = graph.create_named_value("input", DataType::Float32, vec![Dim::Static(2)]);
934 graph.add_input(input);
935 let output = graph.create_named_value("output", DataType::Float32, Shape::new());
936 graph.mark_value_shape_unknown(output);
937 graph.insert_node(Node::new(
938 NodeId(0),
939 "Unsupported",
940 vec![Some(input)],
941 vec![output],
942 ));
943 graph.add_output(output);
944 graph
945 }
946
947 let mut graph = Graph::new();
948 graph.opset_imports.insert(String::new(), 21);
949 let condition = graph.create_named_value("condition", DataType::Bool, Vec::<Dim>::new());
950 graph.add_input(condition);
951 let output = graph.create_named_value("output", DataType::Float32, Shape::new());
952 graph.mark_value_shape_unknown(output);
953 let node = graph.insert_node(Node::new(
954 NodeId(0),
955 "If",
956 vec![Some(condition)],
957 vec![output],
958 ));
959 graph
960 .subgraphs
961 .insert((node, "then_branch".into()), unresolved_branch());
962 graph
963 .subgraphs
964 .insert((node, "else_branch".into()), unresolved_branch());
965 graph.add_output(output);
966
967 let mut model = Model::new(graph);
968 let result = infer_shapes(&mut model).unwrap();
969 assert_eq!(result.unknown, 1);
970 }
971
972 #[test]
973 fn if_inference_preserves_known_scalar_branch_outputs() {
974 fn scalar_branch() -> Graph {
975 let mut graph = Graph::new();
976 let input = graph.create_named_value("input", DataType::Float32, vec![Dim::Static(2)]);
977 graph.add_input(input);
978 let output = graph.create_named_value("output", DataType::Float32, Shape::new());
979 graph.insert_node(Node::new(
980 NodeId(0),
981 "Unsupported",
982 vec![Some(input)],
983 vec![output],
984 ));
985 graph.add_output(output);
986 graph
987 }
988
989 let mut graph = Graph::new();
990 graph.opset_imports.insert(String::new(), 21);
991 let condition = graph.create_named_value("condition", DataType::Bool, Shape::new());
992 graph.add_input(condition);
993 let output = graph.create_named_value("output", DataType::Float32, Shape::new());
994 graph.mark_value_shape_unknown(output);
995 let node = graph.insert_node(Node::new(
996 NodeId(0),
997 "If",
998 vec![Some(condition)],
999 vec![output],
1000 ));
1001 graph
1002 .subgraphs
1003 .insert((node, "then_branch".into()), scalar_branch());
1004 graph
1005 .subgraphs
1006 .insert((node, "else_branch".into()), scalar_branch());
1007 graph.add_output(output);
1008
1009 let mut model = Model::new(graph);
1010 let result = infer_shapes(&mut model).unwrap();
1011 assert!(model.graph.value(output).shape.is_empty());
1012 assert_eq!(result.unknown, 0);
1013 }
1014}