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, 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    merge_traced_leaf_metas, 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::from_typed::<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: first.iter().copied().collect(),
169                    rhs: shape.iter().copied().collect(),
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::CpuEngine;
481    /// use tenferro_runtime::{ExecutionEngine, 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 mut engine = CpuEngine::serial();
491    /// let mut builder = Runtime::builder();
492    /// builder
493    ///     .register_engine(tenferro_cpu::runtime_engine_registration(&engine).unwrap())
494    ///     .unwrap();
495    /// let runtime = builder.build().unwrap();
496    /// let prepared = runtime
497    ///     .prepare(&program, &[], &[engine.target()])
498    ///     .unwrap();
499    /// let mut outputs = runtime
500    ///     .run(&mut [&mut engine], &prepared, &[])
501    ///     .unwrap();
502    /// let out = outputs.remove(0).into_tensor(&mut engine).unwrap();
503    ///
504    /// assert_eq!(
505    ///     out.as_slice::<f64>().unwrap(),
506    ///     &[30.0, 10.0],
507    /// );
508    /// ```
509    ///
510    /// # Errors
511    ///
512    /// Returns [`Error::Validation`] with `InvalidArgument` when the tensor
513    /// shape is not concrete, `AxisOutOfBounds` when `axis` is outside its
514    /// rank, or `InvalidArgument` when a position is outside the selected
515    /// axis extent.
516    pub fn index_select(&self, axis: isize, positions: &[usize]) -> Result<Self> {
517        let shape = try_concrete_shape(self).ok_or_else(|| {
518            Error::validation(
519                "index_select",
520                ErrorPhase::GraphBuild,
521                ValidationError::InvalidArgument {
522                    argument: "shape",
523                    message: "index_select requires a concrete shape hint".into(),
524                },
525            )
526        })?;
527        let (indices_tensor, config, out_shape) = index_select_config(&shape, axis, positions)?;
528        let indices = TracedTensor::from_tensor_concrete_shape(indices_tensor)?;
529        apply_binary_preserve_input_dtypes(
530            StdTensorOp::Gather(config),
531            self,
532            &indices,
533            out_shape.len(),
534            Some(out_shape.into_iter().map(SymDim::from).collect()),
535            self.dtype,
536        )
537    }
538
539    /// Stack tensors along a newly inserted axis.
540    ///
541    /// # Examples
542    ///
543    /// ```
544    /// use tenferro_cpu::CpuEngine;
545    /// use tenferro_runtime::{ExecutionEngine, GraphCompiler, Runtime, Tensor, TracedTensor};
546    ///
547    /// let a = TracedTensor::from_tensor_concrete_shape(Tensor::from_vec_col_major(vec![], vec![1.0_f64]).unwrap()).unwrap();
548    /// let b = TracedTensor::from_tensor_concrete_shape(Tensor::from_vec_col_major(vec![], vec![2.0_f64]).unwrap()).unwrap();
549    /// let stacked = TracedTensor::stack(&[&a, &b], -1).unwrap();
550    /// let mut compiler = GraphCompiler::new();
551    /// let program = compiler.compile(&stacked).unwrap();
552    /// let mut engine = CpuEngine::serial();
553    /// let mut builder = Runtime::builder();
554    /// builder
555    ///     .register_engine(tenferro_cpu::runtime_engine_registration(&engine).unwrap())
556    ///     .unwrap();
557    /// let runtime = builder.build().unwrap();
558    /// let prepared = runtime
559    ///     .prepare(&program, &[], &[engine.target()])
560    ///     .unwrap();
561    /// let mut outputs = runtime
562    ///     .run(&mut [&mut engine], &prepared, &[])
563    ///     .unwrap();
564    /// let out = outputs.remove(0).into_tensor(&mut engine).unwrap();
565    ///
566    /// assert_eq!(
567    ///     out.as_slice::<f64>().unwrap(),
568    ///     &[1.0, 2.0],
569    /// );
570    /// ```
571    ///
572    /// # Errors
573    ///
574    /// Returns [`Error::Validation`] with `InvalidArgument` for an empty input
575    /// list, `ShapeMismatch` for incompatible input shapes, or
576    /// `AxisOutOfBounds` when `dim` is outside the output rank.
577    pub fn stack(tensors: &[&Self], dim: isize) -> Result<Self> {
578        let first = tensors.first().copied().ok_or_else(|| {
579            Error::validation(
580                "stack",
581                ErrorPhase::GraphBuild,
582                ValidationError::InvalidArgument {
583                    argument: "tensors",
584                    message: "stack requires at least one input".into(),
585                },
586            )
587        })?;
588        let first_shape = try_concrete_shape(first).ok_or_else(|| {
589            Error::validation(
590                "stack",
591                ErrorPhase::GraphBuild,
592                ValidationError::InvalidArgument {
593                    argument: "shape",
594                    message: "stack requires concrete shape hints".into(),
595                },
596            )
597        })?;
598        let mut shapes = Vec::with_capacity(tensors.len());
599        shapes.push(first_shape.as_slice());
600        let mut owned_shapes = Vec::with_capacity(tensors.len().saturating_sub(1));
601        for tensor in tensors.iter().copied().skip(1) {
602            owned_shapes.push(try_concrete_shape(tensor).ok_or_else(|| {
603                Error::validation(
604                    "stack",
605                    ErrorPhase::GraphBuild,
606                    ValidationError::InvalidArgument {
607                        argument: "shape",
608                        message: "stack requires concrete shape hints".into(),
609                    },
610                )
611            })?);
612        }
613        shapes.extend(owned_shapes.iter().map(Vec::as_slice));
614        validate_stack_shapes("stack", &shapes)?;
615
616        let axis = normalize_insert_axis("stack", dim, first.rank)?;
617        let mut expanded_shape = first_shape;
618        expanded_shape.insert(axis, 1);
619        let mut out_shape = expanded_shape.clone();
620        out_shape[axis] = tensors.len();
621        let expanded = tensors
622            .iter()
623            .map(|tensor| tensor.reshape(&expanded_shape))
624            .collect::<Result<Vec<_>>>()?;
625        let refs = expanded.iter().collect::<Vec<_>>();
626        apply_nary_concatenate(
627            &refs,
628            axis,
629            out_shape.into_iter().map(SymDim::from).collect(),
630        )
631    }
632
633    /// Concatenate tensors along one existing axis.
634    ///
635    /// # Errors
636    ///
637    /// Returns [`Error::Validation`] with `InvalidArgument` for an empty input
638    /// list, `RankMismatch`/`ShapeMismatch` for incompatible input shapes, or
639    /// `AxisOutOfBounds` when `axis` is outside the input rank.
640    pub fn concatenate(tensors: &[&Self], axis: usize) -> Result<Self> {
641        let first = tensors.first().copied().ok_or_else(|| {
642            Error::validation(
643                "concatenate",
644                ErrorPhase::GraphBuild,
645                ValidationError::InvalidArgument {
646                    argument: "tensors",
647                    message: "concatenate requires at least one input".into(),
648                },
649            )
650        })?;
651        if axis >= first.rank {
652            return Err(Error::validation(
653                "concatenate",
654                ErrorPhase::GraphBuild,
655                ValidationError::AxisOutOfBounds {
656                    axis,
657                    rank: first.rank,
658                },
659            ));
660        }
661        for tensor in tensors.iter().copied().skip(1) {
662            if tensor.rank != first.rank {
663                return Err(Error::validation(
664                    "concatenate",
665                    ErrorPhase::GraphBuild,
666                    ValidationError::RankMismatch {
667                        expected: first.rank,
668                        actual: tensor.rank,
669                    },
670                ));
671            }
672        }
673
674        let op = StdTensorOp::Concatenate {
675            axis,
676            input_count: tensors.len(),
677        };
678        let (_, out_shape_hint) =
679            infer_traced_single_output_shape("TracedTensor::concatenate", &op, tensors)?;
680        let out_shape = out_shape_hint.ok_or_else(|| {
681            Error::Internal("concatenate shape inference returned no shape hint".into())
682        })?;
683        apply_nary_concatenate(tensors, axis, out_shape)
684    }
685}
686
687fn apply_nary_concatenate(
688    tensors: &[&TracedTensor],
689    axis: usize,
690    out_shape: Vec<SymDim>,
691) -> Result<TracedTensor> {
692    let out_dtype = promote_dtypes(tensors.iter().map(|tensor| tensor.dtype));
693    let tensors = tensors
694        .iter()
695        .map(|tensor| {
696            if tensor.dtype != out_dtype {
697                tensor.cast(out_dtype)
698            } else {
699                Ok((*tensor).clone())
700            }
701        })
702        .collect::<Result<Vec<_>>>()?;
703
704    let mut builder = GraphBuilder::new();
705    for tensor in &tensors {
706        builder.add_parent(tensor.graph.clone());
707    }
708    let input_refs = tensors
709        .iter()
710        .map(|tensor| ValueRef::External(tensor.graph.values()[tensor.val].key.clone()))
711        .collect::<Vec<_>>();
712    let outputs = builder.add_operation(
713        StdTensorOp::Concatenate {
714            axis,
715            input_count: tensors.len(),
716        },
717        input_refs,
718        OperationRole::Primary,
719    );
720    builder.set_outputs(outputs.clone());
721    let graph = Arc::new(builder.build());
722    // Callers route through shape inference before graph construction.
723    let metadata_scope =
724        super::traced::register_metadata_or_runtime_state(register_scoped_value_metadata(
725            graph.values()[outputs[0]].key.clone(),
726            tensor_meta(out_dtype, out_shape.clone()),
727        ))?;
728
729    let inputs_map = merge_traced_inputs_map(tensors.iter());
730    let leaf_metas = merge_traced_leaf_metas(tensors.iter());
731    let mut extra_roots = Vec::new();
732    let mut checkpoint_chain = None;
733    for tensor in &tensors {
734        extra_roots.extend(tensor.extra_roots.iter().cloned());
735        checkpoint_chain =
736            CheckpointNode::merge_chains(checkpoint_chain, tensor.checkpoint_chain.clone());
737    }
738    Ok(TracedTensor {
739        id: next_traced_id(),
740        rank: out_shape.len(),
741        dtype: out_dtype,
742        graph,
743        val: outputs[0],
744        data: None,
745        shape_hint: Some(out_shape),
746        inputs_map,
747        leaf_metas,
748        extra_roots,
749        checkpoint_chain,
750        metadata_scopes: MetadataScopeChain::with_new(
751            metadata_scope,
752            tensors.iter().map(|tensor| &tensor.metadata_scopes),
753        ),
754        constraint_scopes: ConstraintScopeChain::merge(
755            tensors.iter().map(|tensor| &tensor.constraint_scopes),
756        ),
757    })
758}
759
760#[cfg(test)]
761mod tests;
tenferro_runtime/shape_packing.rs

tenferro_runtime/
shape_packing.rs
