tenferro_runtime/
shape_packing.rs

1use std::ops::Range;
2use std::sync::Arc;
3
4use computegraph::graph::GraphBuilder;
5use computegraph::types::{OperationRole, ValueRef};
6use tenferro_ops::std_tensor_op::StdTensorOp;
7use tenferro_tensor::{
8    GatherConfig, ShapeMismatch, ShapeVec, SliceConfig, Tensor, TypedTensor, ValidationError,
9};
10
11use crate::checkpoint::CheckpointNode;
12use crate::error::{Error, ErrorPhase, Result};
13use crate::metadata::{register_scoped_value_metadata, tensor_meta, MetadataScopeChain};
14use crate::shape_constraint::ConstraintScopeChain;
15use crate::shape_infer::promote_dtypes;
16use crate::sym_dim::SymDim;
17use crate::traced::{
18    apply_binary_preserve_input_dtypes, infer_traced_single_output_shape, merge_traced_inputs_map,
19    next_traced_id, try_concrete_shape,
20};
21use crate::TracedTensor;
22
23fn normalize_existing_axis(op: &'static str, axis: isize, rank: usize) -> Result<usize> {
24    let normalized = if axis >= 0 {
25        axis as usize
26    } else {
27        rank.checked_sub(axis.unsigned_abs()).ok_or_else(|| {
28            Error::validation(
29                op,
30                ErrorPhase::GraphBuild,
31                ValidationError::AxisOutOfBounds {
32                    axis: axis.unsigned_abs(),
33                    rank,
34                },
35            )
36        })?
37    };
38    if normalized >= rank {
39        return Err(Error::validation(
40            op,
41            ErrorPhase::GraphBuild,
42            ValidationError::AxisOutOfBounds {
43                axis: axis.unsigned_abs(),
44                rank,
45            },
46        ));
47    }
48    Ok(normalized)
49}
50
51fn normalize_insert_axis(op: &'static str, axis: isize, rank: usize) -> Result<usize> {
52    let insert_rank = rank.checked_add(1).ok_or_else(|| {
53        Error::validation(
54            op,
55            ErrorPhase::GraphBuild,
56            ValidationError::AxisOutOfBounds {
57                axis: axis.unsigned_abs(),
58                rank,
59            },
60        )
61    })?;
62    let normalized = if axis >= 0 {
63        axis as usize
64    } else {
65        insert_rank
66            .checked_sub(axis.unsigned_abs())
67            .ok_or_else(|| {
68                Error::validation(
69                    op,
70                    ErrorPhase::GraphBuild,
71                    ValidationError::AxisOutOfBounds {
72                        axis: axis.unsigned_abs(),
73                        rank: insert_rank,
74                    },
75                )
76            })?
77    };
78    if normalized > rank {
79        return Err(Error::validation(
80            op,
81            ErrorPhase::GraphBuild,
82            ValidationError::AxisOutOfBounds {
83                axis: axis.unsigned_abs(),
84                rank: insert_rank,
85            },
86        ));
87    }
88    Ok(normalized)
89}
90
91fn index_select_config(
92    shape: &[usize],
93    axis: isize,
94    positions: &[usize],
95) -> Result<(Tensor, GatherConfig, Vec<usize>)> {
96    let axis = normalize_existing_axis("index_select", axis, shape.len())?;
97    let axis_extent = shape[axis];
98    for &position in positions {
99        if position >= axis_extent {
100            return Err(Error::validation(
101                "index_select",
102                ErrorPhase::GraphBuild,
103                ValidationError::InvalidArgument {
104                    argument: "positions",
105                    message: format!(
106                    "position {position} out of bounds for axis {axis} with extent {axis_extent}"
107                    ),
108                },
109            ));
110        }
111    }
112
113    let mut out_shape = shape.to_vec();
114    out_shape[axis] = positions.len();
115
116    let mut slice_sizes = shape.to_vec();
117    slice_sizes[axis] = 1;
118
119    let offset_dims = (0..shape.len()).filter(|&dim| dim != axis).collect();
120    let index_data = positions
121        .iter()
122        .map(|&position| {
123            i64::try_from(position).map_err(|_| {
124                Error::validation(
125                    "index_select",
126                    ErrorPhase::GraphBuild,
127                    ValidationError::InvalidArgument {
128                        argument: "positions",
129                        message: format!("position {position} cannot be represented as i64"),
130                    },
131                )
132            })
133        })
134        .collect::<Result<Vec<_>>>()?;
135    let indices = Tensor::I64(TypedTensor::from_vec_col_major(
136        vec![positions.len(), 1],
137        index_data,
138    )?);
139
140    let config = GatherConfig {
141        offset_dims,
142        collapsed_slice_dims: vec![axis],
143        start_index_map: vec![axis],
144        index_vector_dim: 1,
145        slice_sizes,
146    };
147
148    Ok((indices, config, out_shape))
149}
150
151fn validate_stack_shapes(op: &'static str, shapes: &[&[usize]]) -> Result<()> {
152    let Some(first) = shapes.first() else {
153        return Err(Error::validation(
154            op,
155            ErrorPhase::GraphBuild,
156            ValidationError::InvalidArgument {
157                argument: "tensors",
158                message: "stack requires at least one input".into(),
159            },
160        ));
161    };
162    for shape in shapes.iter().skip(1) {
163        if *shape != *first {
164            return Err(Error::validation(
165                op,
166                ErrorPhase::GraphBuild,
167                ShapeMismatch::IncompatibleShapes {
168                    lhs: ShapeVec::from_vec(first.to_vec()),
169                    rhs: ShapeVec::from_vec(shape.to_vec()),
170                }
171                .into(),
172            ));
173        }
174    }
175    Ok(())
176}
177
178#[derive(Clone, Debug)]
179enum AxisSelection {
180    Slice {
181        axis: usize,
182        range: Range<usize>,
183        step: usize,
184    },
185    Take {
186        axis: usize,
187        indices: Vec<usize>,
188    },
189}
190
191fn concrete_shape_for_axis_slice(tensor: &TracedTensor, op: &'static str) -> Result<Vec<usize>> {
192    try_concrete_shape(tensor).ok_or_else(|| {
193        Error::validation(
194            op,
195            ErrorPhase::GraphBuild,
196            ValidationError::InvalidArgument {
197                argument: "shape",
198                message: format!("{op} requires a concrete shape hint"),
199            },
200        )
201    })
202}
203
204fn validate_axis_selection(
205    op: &'static str,
206    rank: usize,
207    seen: &mut [bool],
208    axis: usize,
209) -> Result<()> {
210    if axis >= rank {
211        return Err(Error::validation(
212            op,
213            ErrorPhase::GraphBuild,
214            ValidationError::AxisOutOfBounds { axis, rank },
215        ));
216    }
217    if seen[axis] {
218        return Err(Error::validation(
219            op,
220            ErrorPhase::GraphBuild,
221            ValidationError::DuplicateAxis {
222                axis,
223                role: "selection",
224            },
225        ));
226    }
227    seen[axis] = true;
228    Ok(())
229}
230
231fn apply_slice_axis_config(
232    op: &'static str,
233    shape: &[usize],
234    selections: &[AxisSelection],
235) -> Result<Option<SliceConfig>> {
236    let mut starts = vec![0; shape.len()];
237    let mut limits = shape.to_vec();
238    let mut strides = vec![1; shape.len()];
239    let mut has_slice = false;
240    for selection in selections {
241        let AxisSelection::Slice { axis, range, step } = selection else {
242            continue;
243        };
244        if *step == 0 {
245            return Err(Error::validation(
246                op,
247                ErrorPhase::GraphBuild,
248                ValidationError::InvalidSliceStep { step: 0 },
249            ));
250        }
251        let extent = shape[*axis];
252        if range.start > range.end || range.end > extent {
253            let start = isize::try_from(range.start).map_err(|_| {
254                Error::validation(op, ErrorPhase::GraphBuild, ValidationError::IntegerOverflow)
255            })?;
256            let end = isize::try_from(range.end).map_err(|_| {
257                Error::validation(op, ErrorPhase::GraphBuild, ValidationError::IntegerOverflow)
258            })?;
259            return Err(Error::validation(
260                op,
261                ErrorPhase::GraphBuild,
262                ValidationError::InvalidSliceBounds {
263                    start,
264                    end,
265                    axis_len: extent,
266                },
267            ));
268        }
269        starts[*axis] = range.start;
270        limits[*axis] = range.end;
271        strides[*axis] = *step;
272        has_slice = true;
273    }
274    Ok(has_slice.then_some(SliceConfig {
275        starts,
276        limits,
277        strides,
278    }))
279}
280
281/// Rank-preserving traced tensor slicing builder.
282///
283/// Unspecified axes are kept whole. Range selections become one `Slice`
284/// operation; host-known position selections become `Gather`/`index_select`
285/// operations.
286///
287/// # Examples
288///
289/// ```rust
290/// use tenferro_runtime::TracedTensor;
291///
292/// let x = TracedTensor::from_vec_col_major(vec![3, 4], vec![0.0_f64; 12]).unwrap();
293/// let y = x.slice_builder().axis(0, 0..2).axis_step(1, 0..4, 2).apply().unwrap();
294/// assert_eq!(y.try_concrete_shape(), Some(vec![2, 2]));
295/// ```
296#[derive(Clone, Debug)]
297pub struct TracedSliceBuilder<'a> {
298    tensor: &'a TracedTensor,
299    selections: Vec<AxisSelection>,
300}
301
302impl<'a> TracedSliceBuilder<'a> {
303    fn new(tensor: &'a TracedTensor) -> Self {
304        Self {
305            tensor,
306            selections: Vec::new(),
307        }
308    }
309
310    /// Add an exclusive-end range selection for one axis.
311    ///
312    /// # Examples
313    ///
314    /// ```rust
315    /// use tenferro_runtime::TracedTensor;
316    ///
317    /// let x = TracedTensor::from_vec_col_major(vec![4], vec![1.0_f64, 2.0, 3.0, 4.0]).unwrap();
318    /// let y = x.slice_builder().axis(0, 1..3).apply().unwrap();
319    /// assert_eq!(y.try_concrete_shape(), Some(vec![2]));
320    /// ```
321    pub fn axis(mut self, axis: usize, range: Range<usize>) -> Self {
322        self.selections.push(AxisSelection::Slice {
323            axis,
324            range,
325            step: 1,
326        });
327        self
328    }
329
330    /// Add an exclusive-end strided range selection for one axis.
331    ///
332    /// # Examples
333    ///
334    /// ```rust
335    /// use tenferro_runtime::TracedTensor;
336    ///
337    /// let x = TracedTensor::from_vec_col_major(vec![5], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0]).unwrap();
338    /// let y = x.slice_builder().axis_step(0, 0..5, 2).apply().unwrap();
339    /// assert_eq!(y.try_concrete_shape(), Some(vec![3]));
340    /// ```
341    pub fn axis_step(mut self, axis: usize, range: Range<usize>, step: usize) -> Self {
342        self.selections
343            .push(AxisSelection::Slice { axis, range, step });
344        self
345    }
346
347    /// Add a host-known position selection for one axis.
348    ///
349    /// # Examples
350    ///
351    /// ```rust
352    /// use tenferro_runtime::TracedTensor;
353    ///
354    /// let x = TracedTensor::from_vec_col_major(vec![3], vec![1.0_f64, 2.0, 3.0]).unwrap();
355    /// let y = x.slice_builder().take_axis(0, &[2, 0]).apply().unwrap();
356    /// assert_eq!(y.try_concrete_shape(), Some(vec![2]));
357    /// ```
358    pub fn take_axis(mut self, axis: usize, indices: &[usize]) -> Self {
359        self.selections.push(AxisSelection::Take {
360            axis,
361            indices: indices.to_vec(),
362        });
363        self
364    }
365
366    /// Build and apply the requested slice/take operations.
367    ///
368    /// # Examples
369    ///
370    /// ```rust
371    /// use tenferro_runtime::TracedTensor;
372    ///
373    /// let x = TracedTensor::from_vec_col_major(vec![4], vec![1.0_f64, 2.0, 3.0, 4.0]).unwrap();
374    /// let y = x.slice_builder().axis(0, 1..4).apply().unwrap();
375    /// assert_eq!(y.try_concrete_shape(), Some(vec![3]));
376    /// ```
377    ///
378    /// # Errors
379    ///
380    /// Returns [`Error::Validation`] with `AxisOutOfBounds` or
381    /// `DuplicateAxis` when selections are invalid, and propagates the
382    /// underlying [`Error::Validation`] from a slice/take graph operation
383    /// that cannot be built.
384    pub fn apply(self) -> Result<TracedTensor> {
385        let shape = concrete_shape_for_axis_slice(self.tensor, "slice_builder")?;
386        let mut seen = vec![false; shape.len()];
387        for selection in &self.selections {
388            let axis = match selection {
389                AxisSelection::Slice { axis, .. } | AxisSelection::Take { axis, .. } => *axis,
390            };
391            validate_axis_selection("slice_builder", shape.len(), &mut seen, axis)?;
392        }
393
394        let mut output = self.tensor.clone();
395        if let Some(config) = apply_slice_axis_config("slice_builder", &shape, &self.selections)? {
396            output = output.slice(config)?;
397        }
398        for selection in self.selections {
399            if let AxisSelection::Take { axis, indices } = selection {
400                output = output.take_axis(axis, &indices)?;
401            }
402        }
403        Ok(output)
404    }
405}
406
407impl TracedTensor {
408    /// Slice one axis with an exclusive-end range, keeping all other axes.
409    ///
410    /// # Examples
411    ///
412    /// ```rust
413    /// use tenferro_runtime::TracedTensor;
414    ///
415    /// let x = TracedTensor::from_vec_col_major(vec![4], vec![1.0_f64, 2.0, 3.0, 4.0]).unwrap();
416    /// let y = x.slice_axis(0, 1..3).unwrap();
417    /// assert_eq!(y.try_concrete_shape(), Some(vec![2]));
418    /// ```
419    ///
420    /// # Errors
421    ///
422    /// Returns [`Error::Validation`] with `AxisOutOfBounds` when `axis` is
423    /// outside the concrete rank, or `InvalidArgument` when `range` is
424    /// outside the selected axis extent.
425    pub fn slice_axis(&self, axis: usize, range: Range<usize>) -> Result<Self> {
426        self.slice_builder().axis(axis, range).apply()
427    }
428
429    /// Start a rank-preserving slicing builder for this tensor.
430    ///
431    /// # Examples
432    ///
433    /// ```rust
434    /// use tenferro_runtime::TracedTensor;
435    ///
436    /// let x = TracedTensor::from_vec_col_major(vec![3], vec![1.0_f64, 2.0, 3.0]).unwrap();
437    /// let y = x.slice_builder().axis(0, 0..2).apply().unwrap();
438    /// assert_eq!(y.try_concrete_shape(), Some(vec![2]));
439    /// ```
440    pub fn slice_builder(&self) -> TracedSliceBuilder<'_> {
441        TracedSliceBuilder::new(self)
442    }
443
444    /// Select entries from one axis using host-known indices.
445    ///
446    /// # Examples
447    ///
448    /// ```rust
449    /// use tenferro_runtime::TracedTensor;
450    ///
451    /// let x = TracedTensor::from_vec_col_major(vec![3], vec![10.0_f64, 20.0, 30.0]).unwrap();
452    /// let y = x.take_axis(0, &[2, 0]).unwrap();
453    /// assert_eq!(y.try_concrete_shape(), Some(vec![2]));
454    /// ```
455    ///
456    /// # Errors
457    ///
458    /// Returns [`Error::Validation`] with `AxisOutOfBounds` when `axis` is
459    /// outside the concrete rank, or `InvalidArgument` when an index list
460    /// cannot be applied to the selected axis.
461    pub fn take_axis(&self, axis: usize, indices: &[usize]) -> Result<Self> {
462        let axis = isize::try_from(axis).map_err(|_| {
463            Error::validation(
464                "take_axis",
465                ErrorPhase::GraphBuild,
466                ValidationError::InvalidArgument {
467                    argument: "axis",
468                    message: format!("axis {axis} cannot be represented as isize"),
469                },
470            )
471        })?;
472        self.index_select(axis, indices)
473    }
474
475    /// Select entries from one axis using host-known positions.
476    ///
477    /// # Examples
478    ///
479    /// ```
480    /// use tenferro_cpu::CpuBackend;
481    /// use tenferro_runtime::{GraphCompiler, Runtime, Tensor, TracedTensor};
482    ///
483    /// let x = TracedTensor::from_tensor_concrete_shape(
484    ///     Tensor::from_vec_col_major(vec![3], vec![10.0_f64, 20.0, 30.0]).unwrap(),
485    /// )
486    /// .unwrap();
487    /// let y = x.index_select(-1, &[2, 0]).unwrap();
488    /// let mut compiler = GraphCompiler::new();
489    /// let program = compiler.compile(&y).unwrap();
490    /// let backend = CpuBackend::new();
491    /// let mut builder = Runtime::builder();
492    /// builder
493    ///     .register_engine(tenferro_cpu::runtime_engine_registration(&backend).unwrap())
494    ///     .unwrap();
495    /// let runtime = builder.build().unwrap();
496    /// let outputs = runtime.run_compiled(&program, &[]).unwrap();
497    /// let out = &outputs[0];
498    ///
499    /// assert_eq!(
500    ///     out.as_slice::<f64>().unwrap(),
501    ///     &[30.0, 10.0],
502    /// );
503    /// ```
504    ///
505    /// # Errors
506    ///
507    /// Returns [`Error::Validation`] with `InvalidArgument` when the tensor
508    /// shape is not concrete, `AxisOutOfBounds` when `axis` is outside its
509    /// rank, or `InvalidArgument` when a position is outside the selected
510    /// axis extent.
511    pub fn index_select(&self, axis: isize, positions: &[usize]) -> Result<Self> {
512        let shape = try_concrete_shape(self).ok_or_else(|| {
513            Error::validation(
514                "index_select",
515                ErrorPhase::GraphBuild,
516                ValidationError::InvalidArgument {
517                    argument: "shape",
518                    message: "index_select requires a concrete shape hint".into(),
519                },
520            )
521        })?;
522        let (indices_tensor, config, out_shape) = index_select_config(&shape, axis, positions)?;
523        let indices = TracedTensor::from_tensor_concrete_shape(indices_tensor)?;
524        apply_binary_preserve_input_dtypes(
525            StdTensorOp::Gather(config),
526            self,
527            &indices,
528            out_shape.len(),
529            Some(out_shape.into_iter().map(SymDim::from).collect()),
530            self.dtype,
531        )
532    }
533
534    /// Stack tensors along a newly inserted axis.
535    ///
536    /// # Examples
537    ///
538    /// ```
539    /// use tenferro_cpu::CpuBackend;
540    /// use tenferro_runtime::{GraphCompiler, Runtime, Tensor, TracedTensor};
541    ///
542    /// let a = TracedTensor::from_tensor_concrete_shape(Tensor::from_vec_col_major(vec![], vec![1.0_f64]).unwrap()).unwrap();
543    /// let b = TracedTensor::from_tensor_concrete_shape(Tensor::from_vec_col_major(vec![], vec![2.0_f64]).unwrap()).unwrap();
544    /// let stacked = TracedTensor::stack(&[&a, &b], -1).unwrap();
545    /// let mut compiler = GraphCompiler::new();
546    /// let program = compiler.compile(&stacked).unwrap();
547    /// let backend = CpuBackend::new();
548    /// let mut builder = Runtime::builder();
549    /// builder
550    ///     .register_engine(tenferro_cpu::runtime_engine_registration(&backend).unwrap())
551    ///     .unwrap();
552    /// let runtime = builder.build().unwrap();
553    /// let outputs = runtime.run_compiled(&program, &[]).unwrap();
554    /// let out = &outputs[0];
555    ///
556    /// assert_eq!(
557    ///     out.as_slice::<f64>().unwrap(),
558    ///     &[1.0, 2.0],
559    /// );
560    /// ```
561    ///
562    /// # Errors
563    ///
564    /// Returns [`Error::Validation`] with `InvalidArgument` for an empty input
565    /// list, `ShapeMismatch` for incompatible input shapes, or
566    /// `AxisOutOfBounds` when `dim` is outside the output rank.
567    pub fn stack(tensors: &[&Self], dim: isize) -> Result<Self> {
568        let first = tensors.first().copied().ok_or_else(|| {
569            Error::validation(
570                "stack",
571                ErrorPhase::GraphBuild,
572                ValidationError::InvalidArgument {
573                    argument: "tensors",
574                    message: "stack requires at least one input".into(),
575                },
576            )
577        })?;
578        let first_shape = try_concrete_shape(first).ok_or_else(|| {
579            Error::validation(
580                "stack",
581                ErrorPhase::GraphBuild,
582                ValidationError::InvalidArgument {
583                    argument: "shape",
584                    message: "stack requires concrete shape hints".into(),
585                },
586            )
587        })?;
588        let mut shapes = Vec::with_capacity(tensors.len());
589        shapes.push(first_shape.as_slice());
590        let mut owned_shapes = Vec::with_capacity(tensors.len().saturating_sub(1));
591        for tensor in tensors.iter().copied().skip(1) {
592            owned_shapes.push(try_concrete_shape(tensor).ok_or_else(|| {
593                Error::validation(
594                    "stack",
595                    ErrorPhase::GraphBuild,
596                    ValidationError::InvalidArgument {
597                        argument: "shape",
598                        message: "stack requires concrete shape hints".into(),
599                    },
600                )
601            })?);
602        }
603        shapes.extend(owned_shapes.iter().map(Vec::as_slice));
604        validate_stack_shapes("stack", &shapes)?;
605
606        let axis = normalize_insert_axis("stack", dim, first.rank)?;
607        let mut expanded_shape = first_shape;
608        expanded_shape.insert(axis, 1);
609        let mut out_shape = expanded_shape.clone();
610        out_shape[axis] = tensors.len();
611        let expanded = tensors
612            .iter()
613            .map(|tensor| tensor.reshape(&expanded_shape))
614            .collect::<Result<Vec<_>>>()?;
615        let refs = expanded.iter().collect::<Vec<_>>();
616        apply_nary_concatenate(
617            &refs,
618            axis,
619            out_shape.into_iter().map(SymDim::from).collect(),
620        )
621    }
622
623    /// Concatenate tensors along one existing axis.
624    ///
625    /// # Errors
626    ///
627    /// Returns [`Error::Validation`] with `InvalidArgument` for an empty input
628    /// list, `RankMismatch`/`ShapeMismatch` for incompatible input shapes, or
629    /// `AxisOutOfBounds` when `axis` is outside the input rank.
630    pub fn concatenate(tensors: &[&Self], axis: usize) -> Result<Self> {
631        let first = tensors.first().copied().ok_or_else(|| {
632            Error::validation(
633                "concatenate",
634                ErrorPhase::GraphBuild,
635                ValidationError::InvalidArgument {
636                    argument: "tensors",
637                    message: "concatenate requires at least one input".into(),
638                },
639            )
640        })?;
641        if axis >= first.rank {
642            return Err(Error::validation(
643                "concatenate",
644                ErrorPhase::GraphBuild,
645                ValidationError::AxisOutOfBounds {
646                    axis,
647                    rank: first.rank,
648                },
649            ));
650        }
651        for tensor in tensors.iter().copied().skip(1) {
652            if tensor.rank != first.rank {
653                return Err(Error::validation(
654                    "concatenate",
655                    ErrorPhase::GraphBuild,
656                    ValidationError::RankMismatch {
657                        expected: first.rank,
658                        actual: tensor.rank,
659                    },
660                ));
661            }
662        }
663
664        let op = StdTensorOp::Concatenate {
665            axis,
666            input_count: tensors.len(),
667        };
668        let (_, out_shape_hint) =
669            infer_traced_single_output_shape("TracedTensor::concatenate", &op, tensors)?;
670        let out_shape = out_shape_hint.ok_or_else(|| {
671            Error::Internal("concatenate shape inference returned no shape hint".into())
672        })?;
673        apply_nary_concatenate(tensors, axis, out_shape)
674    }
675}
676
677fn apply_nary_concatenate(
678    tensors: &[&TracedTensor],
679    axis: usize,
680    out_shape: Vec<SymDim>,
681) -> Result<TracedTensor> {
682    let out_dtype = promote_dtypes(tensors.iter().map(|tensor| tensor.dtype));
683    let tensors = tensors
684        .iter()
685        .map(|tensor| {
686            if tensor.dtype != out_dtype {
687                tensor.cast(out_dtype)
688            } else {
689                Ok((*tensor).clone())
690            }
691        })
692        .collect::<Result<Vec<_>>>()?;
693
694    let mut builder = GraphBuilder::new();
695    for tensor in &tensors {
696        builder.add_parent(tensor.graph.clone());
697    }
698    let input_refs = tensors
699        .iter()
700        .map(|tensor| ValueRef::External(tensor.graph.values()[tensor.val].key.clone()))
701        .collect::<Vec<_>>();
702    let outputs = builder.add_operation(
703        StdTensorOp::Concatenate {
704            axis,
705            input_count: tensors.len(),
706        },
707        input_refs,
708        OperationRole::Primary,
709    );
710    builder.set_outputs(outputs.clone());
711    let graph = Arc::new(builder.build());
712    // Callers route through shape inference before graph construction.
713    let metadata_scope =
714        super::traced::register_metadata_or_runtime_state(register_scoped_value_metadata(
715            graph.values()[outputs[0]].key.clone(),
716            tensor_meta(out_dtype, out_shape.clone()),
717        ))?;
718
719    let inputs_map = merge_traced_inputs_map(tensors.iter());
720    let mut extra_roots = Vec::new();
721    let mut checkpoint_chain = None;
722    for tensor in &tensors {
723        extra_roots.extend(tensor.extra_roots.iter().cloned());
724        checkpoint_chain =
725            CheckpointNode::merge_chains(checkpoint_chain, tensor.checkpoint_chain.clone());
726    }
727    Ok(TracedTensor {
728        id: next_traced_id(),
729        rank: out_shape.len(),
730        dtype: out_dtype,
731        graph,
732        val: outputs[0],
733        data: None,
734        shape_hint: Some(out_shape),
735        inputs_map,
736        extra_roots,
737        checkpoint_chain,
738        metadata_scopes: MetadataScopeChain::with_new(
739            metadata_scope,
740            tensors.iter().map(|tensor| &tensor.metadata_scopes),
741        ),
742        constraint_scopes: ConstraintScopeChain::merge(
743            tensors.iter().map(|tensor| &tensor.constraint_scopes),
744        ),
745    })
746}
747
748#[cfg(test)]
749mod tests {
750    use std::ops::Range;
751
752    use tenferro_tensor::{DType, ErrorKind, ValidationError, ValidationKind};
753
754    use super::{
755        apply_slice_axis_config, concrete_shape_for_axis_slice, index_select_config,
756        normalize_existing_axis, normalize_insert_axis, validate_axis_selection,
757        validate_stack_shapes, AxisSelection,
758    };
759    use crate::TracedTensor;
760
761    #[test]
762    fn axis_normalization_handles_ranks_larger_than_isize_max() {
763        assert_eq!(normalize_existing_axis("test", 0, usize::MAX).unwrap(), 0);
764        assert_eq!(
765            normalize_existing_axis("test", -1, usize::MAX).unwrap(),
766            usize::MAX - 1
767        );
768        assert_eq!(
769            normalize_insert_axis("test", -1, usize::MAX - 1).unwrap(),
770            usize::MAX - 1
771        );
772        assert!(normalize_insert_axis("test", -1, usize::MAX).is_err());
773    }
774
775    #[test]
776    fn axis_normalization_reports_positive_and_negative_bounds() {
777        assert_eq!(normalize_existing_axis("test", 1, 3).unwrap(), 1);
778        assert_eq!(normalize_existing_axis("test", -1, 3).unwrap(), 2);
779        assert_eq!(normalize_insert_axis("test", 3, 3).unwrap(), 3);
780        assert_eq!(normalize_insert_axis("test", -1, 3).unwrap(), 3);
781
782        for error in [
783            normalize_existing_axis("test", 3, 3).unwrap_err(),
784            normalize_existing_axis("test", -4, 3).unwrap_err(),
785            normalize_insert_axis("test", 4, 3).unwrap_err(),
786            normalize_insert_axis("test", -5, 3).unwrap_err(),
787        ] {
788            assert_eq!(
789                error.kind(),
790                ErrorKind::Validation(ValidationKind::AxisOutOfBounds)
791            );
792        }
793    }
794
795    #[test]
796    fn index_select_config_validates_positions_and_builds_gather_metadata() {
797        let (indices, config, output_shape) = index_select_config(&[3, 4], -1, &[3, 1]).unwrap();
798        assert_eq!(indices.as_slice::<i64>().unwrap(), &[3, 1]);
799        assert_eq!(output_shape, vec![3, 2]);
800        assert_eq!(config.collapsed_slice_dims, vec![1]);
801        assert_eq!(config.offset_dims, vec![0]);
802        assert_eq!(config.start_index_map, vec![1]);
803        assert_eq!(config.slice_sizes, vec![3, 1]);
804
805        let error = index_select_config(&[3], 0, &[3]).unwrap_err();
806        assert!(matches!(
807            error,
808            crate::Error::Validation {
809                source: ValidationError::InvalidArgument {
810                    argument: "positions",
811                    ..
812                },
813                ..
814            }
815        ));
816    }
817
818    #[test]
819    fn stack_shape_and_selection_helpers_cover_validation_contracts() {
820        let first = [2usize, 3];
821        let same = [2usize, 3];
822        let different = [2usize, 4];
823        assert!(validate_stack_shapes("stack", &[&first, &same]).is_ok());
824        assert!(matches!(
825            validate_stack_shapes("stack", &[]).unwrap_err(),
826            crate::Error::Validation {
827                source: ValidationError::InvalidArgument {
828                    argument: "tensors",
829                    ..
830                },
831                ..
832            }
833        ));
834        assert!(matches!(
835            validate_stack_shapes("stack", &[&first, &different]).unwrap_err(),
836            crate::Error::Validation {
837                source: ValidationError::ShapeMismatch(_),
838                ..
839            }
840        ));
841
842        let mut seen = vec![false; 2];
843        validate_axis_selection("slice", 2, &mut seen, 1).unwrap();
844        assert!(matches!(
845            validate_axis_selection("slice", 2, &mut seen, 1).unwrap_err(),
846            crate::Error::Validation {
847                source: ValidationError::DuplicateAxis { axis: 1, .. },
848                ..
849            }
850        ));
851        assert!(matches!(
852            validate_axis_selection("slice", 2, &mut seen, 2).unwrap_err(),
853            crate::Error::Validation {
854                source: ValidationError::AxisOutOfBounds { axis: 2, rank: 2 },
855                ..
856            }
857        ));
858    }
859
860    #[test]
861    fn slice_config_distinguishes_take_only_from_ranges_and_rejects_bad_ranges() {
862        let take_only = [AxisSelection::Take {
863            axis: 0,
864            indices: vec![1],
865        }];
866        assert!(apply_slice_axis_config("slice", &[3], &take_only)
867            .unwrap()
868            .is_none());
869
870        let ranges = [AxisSelection::Slice {
871            axis: 0,
872            range: 1..3,
873            step: 2,
874        }];
875        assert_eq!(
876            apply_slice_axis_config("slice", &[4], &ranges)
877                .unwrap()
878                .unwrap()
879                .strides,
880            vec![2]
881        );
882
883        let zero_step = [AxisSelection::Slice {
884            axis: 0,
885            range: 0..1,
886            step: 0,
887        }];
888        assert!(matches!(
889            apply_slice_axis_config("slice", &[2], &zero_step).unwrap_err(),
890            crate::Error::Validation {
891                source: ValidationError::InvalidSliceStep { step: 0 },
892                ..
893            }
894        ));
895
896        let bad_bounds = [AxisSelection::Slice {
897            axis: 0,
898            range: Range { start: 2, end: 4 },
899            step: 1,
900        }];
901        assert!(matches!(
902            apply_slice_axis_config("slice", &[3], &bad_bounds).unwrap_err(),
903            crate::Error::Validation {
904                source: ValidationError::InvalidSliceBounds { axis_len: 3, .. },
905                ..
906            }
907        ));
908    }
909
910    #[test]
911    fn public_slice_stack_and_concatenate_paths_preserve_structured_errors() {
912        let x = TracedTensor::from_vec_col_major(
913            vec![4, 5],
914            (0..20).map(|value| value as f64).collect(),
915        )
916        .unwrap();
917        let sliced = x
918            .slice_builder()
919            .axis(0, 1..4)
920            .take_axis(1, &[4, 0])
921            .apply()
922            .unwrap();
923        assert_eq!(sliced.try_concrete_shape(), Some(vec![3, 2]));
924
925        assert!(matches!(
926            x.slice_builder()
927                .axis(0, 0..2)
928                .axis_step(0, 0..2, 1)
929                .apply(),
930            Err(crate::Error::Validation {
931                source: ValidationError::DuplicateAxis { .. },
932                ..
933            })
934        ));
935
936        let symbolic = TracedTensor::input_symbolic_shape(DType::F64, 2).unwrap();
937        assert!(matches!(
938            concrete_shape_for_axis_slice(&symbolic, "slice_builder"),
939            Err(crate::Error::Validation {
940                source: ValidationError::InvalidArgument {
941                    argument: "shape",
942                    ..
943                },
944                ..
945            })
946        ));
947        assert!(matches!(
948            symbolic.index_select(0, &[0]),
949            Err(crate::Error::Validation {
950                source: ValidationError::InvalidArgument {
951                    argument: "shape",
952                    ..
953                },
954                ..
955            })
956        ));
957
958        let a = TracedTensor::from_vec_col_major(vec![2], vec![1.0_f64, 2.0]).unwrap();
959        let b = TracedTensor::from_vec_col_major(vec![2], vec![3.0_f64, 4.0]).unwrap();
960        let stacked = TracedTensor::stack(&[&a, &b], -1).unwrap();
961        assert_eq!(stacked.try_concrete_shape(), Some(vec![2, 2]));
962        assert!(matches!(
963            TracedTensor::stack(&[], 0),
964            Err(crate::Error::Validation {
965                source: ValidationError::InvalidArgument {
966                    argument: "tensors",
967                    ..
968                },
969                ..
970            })
971        ));
972        assert!(matches!(
973            TracedTensor::stack(&[&a, &x], 0),
974            Err(crate::Error::Validation {
975                source: ValidationError::ShapeMismatch(_),
976                ..
977            })
978        ));
979        assert!(matches!(
980            TracedTensor::stack(&[&a, &b], 2),
981            Err(crate::Error::Validation {
982                source: ValidationError::AxisOutOfBounds { .. },
983                ..
984            })
985        ));
986
987        let c = TracedTensor::from_vec_col_major(vec![2], vec![3.0_f32, 4.0]).unwrap();
988        let concatenated = TracedTensor::concatenate(&[&a, &c], 0).unwrap();
989        assert_eq!(concatenated.dtype, DType::F64);
990        assert_eq!(concatenated.try_concrete_shape(), Some(vec![4]));
991        assert!(matches!(
992            TracedTensor::concatenate(&[], 0),
993            Err(crate::Error::Validation {
994                source: ValidationError::InvalidArgument {
995                    argument: "tensors",
996                    ..
997                },
998                ..
999            })
1000        ));
1001        assert!(matches!(
1002            TracedTensor::concatenate(&[&a, &b], 1),
1003            Err(crate::Error::Validation {
1004                source: ValidationError::AxisOutOfBounds { .. },
1005                ..
1006            })
1007        ));
1008        let rank_one = TracedTensor::from_vec_col_major(vec![1, 2], vec![1.0_f64, 2.0]).unwrap();
1009        assert!(matches!(
1010            TracedTensor::concatenate(&[&a, &rank_one], 0),
1011            Err(crate::Error::Validation {
1012                source: ValidationError::RankMismatch { .. },
1013                ..
1014            })
1015        ));
1016    }
1017}
tenferro_runtime/shape_packing.rs

tenferro_runtime/
shape_packing.rs
