tenferro_runtime/
shape_packing.rs

1use std::ops::Range;
2
3use tenferro_ops::std_tensor_op::StdTensorOp;
4use tenferro_tensor::{
5    GatherConfig, ShapeMismatch, SliceConfig, Tensor, TypedTensor, ValidationError,
6};
7
8use crate::error::{Error, ErrorPhase, Result};
9use crate::shape_infer::promote_dtypes;
10use crate::sym_dim::SymDim;
11use crate::traced::{
12    apply_binary_preserve_input_dtypes, infer_traced_single_output_shape, try_concrete_shape,
13};
14use crate::TracedTensor;
15
16fn normalize_existing_axis(op: &'static str, axis: isize, rank: usize) -> Result<usize> {
17    let normalized = if axis >= 0 {
18        axis as usize
19    } else {
20        rank.checked_sub(axis.unsigned_abs()).ok_or_else(|| {
21            Error::validation(
22                op,
23                ErrorPhase::GraphBuild,
24                ValidationError::AxisOutOfBounds {
25                    axis: axis.unsigned_abs(),
26                    rank,
27                },
28            )
29        })?
30    };
31    if normalized >= rank {
32        return Err(Error::validation(
33            op,
34            ErrorPhase::GraphBuild,
35            ValidationError::AxisOutOfBounds {
36                axis: axis.unsigned_abs(),
37                rank,
38            },
39        ));
40    }
41    Ok(normalized)
42}
43
44fn normalize_insert_axis(op: &'static str, axis: isize, rank: usize) -> Result<usize> {
45    let insert_rank = rank.checked_add(1).ok_or_else(|| {
46        Error::validation(
47            op,
48            ErrorPhase::GraphBuild,
49            ValidationError::AxisOutOfBounds {
50                axis: axis.unsigned_abs(),
51                rank,
52            },
53        )
54    })?;
55    let normalized = if axis >= 0 {
56        axis as usize
57    } else {
58        insert_rank
59            .checked_sub(axis.unsigned_abs())
60            .ok_or_else(|| {
61                Error::validation(
62                    op,
63                    ErrorPhase::GraphBuild,
64                    ValidationError::AxisOutOfBounds {
65                        axis: axis.unsigned_abs(),
66                        rank: insert_rank,
67                    },
68                )
69            })?
70    };
71    if normalized > rank {
72        return Err(Error::validation(
73            op,
74            ErrorPhase::GraphBuild,
75            ValidationError::AxisOutOfBounds {
76                axis: axis.unsigned_abs(),
77                rank: insert_rank,
78            },
79        ));
80    }
81    Ok(normalized)
82}
83
84fn index_select_config(
85    shape: &[usize],
86    axis: isize,
87    positions: &[usize],
88) -> Result<(Tensor, GatherConfig, Vec<usize>)> {
89    let axis = normalize_existing_axis("index_select", axis, shape.len())?;
90    let axis_extent = shape[axis];
91    for &position in positions {
92        if position >= axis_extent {
93            return Err(Error::validation(
94                "index_select",
95                ErrorPhase::GraphBuild,
96                ValidationError::InvalidArgument {
97                    argument: "positions",
98                    message: format!(
99                    "position {position} out of bounds for axis {axis} with extent {axis_extent}"
100                    ),
101                },
102            ));
103        }
104    }
105
106    let mut out_shape = shape.to_vec();
107    out_shape[axis] = positions.len();
108
109    let mut slice_sizes = shape.to_vec();
110    slice_sizes[axis] = 1;
111
112    let offset_dims = (0..shape.len()).filter(|&dim| dim != axis).collect();
113    let index_data = positions
114        .iter()
115        .map(|&position| {
116            i64::try_from(position).map_err(|_| {
117                Error::validation(
118                    "index_select",
119                    ErrorPhase::GraphBuild,
120                    ValidationError::InvalidArgument {
121                        argument: "positions",
122                        message: format!("position {position} cannot be represented as i64"),
123                    },
124                )
125            })
126        })
127        .collect::<Result<Vec<_>>>()?;
128    let indices = Tensor::from_typed::<i64>(TypedTensor::from_vec_col_major(
129        vec![positions.len(), 1],
130        index_data,
131    )?);
132
133    let config = GatherConfig {
134        offset_dims,
135        collapsed_slice_dims: vec![axis],
136        start_index_map: vec![axis],
137        index_vector_dim: 1,
138        slice_sizes,
139    };
140
141    Ok((indices, config, out_shape))
142}
143
144fn validate_stack_shapes(op: &'static str, shapes: &[&[usize]]) -> Result<()> {
145    let Some(first) = shapes.first() else {
146        return Err(Error::validation(
147            op,
148            ErrorPhase::GraphBuild,
149            ValidationError::InvalidArgument {
150                argument: "tensors",
151                message: "stack requires at least one input".into(),
152            },
153        ));
154    };
155    for shape in shapes.iter().skip(1) {
156        if *shape != *first {
157            return Err(Error::validation(
158                op,
159                ErrorPhase::GraphBuild,
160                ShapeMismatch::IncompatibleShapes {
161                    lhs: first.iter().copied().collect(),
162                    rhs: shape.iter().copied().collect(),
163                }
164                .into(),
165            ));
166        }
167    }
168    Ok(())
169}
170
171#[derive(Clone, Debug)]
172enum AxisSelection {
173    Slice {
174        axis: usize,
175        range: Range<usize>,
176        step: usize,
177    },
178    Take {
179        axis: usize,
180        indices: Vec<usize>,
181    },
182}
183
184fn concrete_shape_for_axis_slice(tensor: &TracedTensor, op: &'static str) -> Result<Vec<usize>> {
185    try_concrete_shape(tensor).ok_or_else(|| {
186        Error::validation(
187            op,
188            ErrorPhase::GraphBuild,
189            ValidationError::InvalidArgument {
190                argument: "shape",
191                message: format!("{op} requires a concrete shape hint"),
192            },
193        )
194    })
195}
196
197fn validate_axis_selection(
198    op: &'static str,
199    rank: usize,
200    seen: &mut [bool],
201    axis: usize,
202) -> Result<()> {
203    if axis >= rank {
204        return Err(Error::validation(
205            op,
206            ErrorPhase::GraphBuild,
207            ValidationError::AxisOutOfBounds { axis, rank },
208        ));
209    }
210    if seen[axis] {
211        return Err(Error::validation(
212            op,
213            ErrorPhase::GraphBuild,
214            ValidationError::DuplicateAxis {
215                axis,
216                role: "selection",
217            },
218        ));
219    }
220    seen[axis] = true;
221    Ok(())
222}
223
224fn apply_slice_axis_config(
225    op: &'static str,
226    shape: &[usize],
227    selections: &[AxisSelection],
228) -> Result<Option<SliceConfig>> {
229    let mut starts = vec![0; shape.len()];
230    let mut limits = shape.to_vec();
231    let mut strides = vec![1; shape.len()];
232    let mut has_slice = false;
233    for selection in selections {
234        let AxisSelection::Slice { axis, range, step } = selection else {
235            continue;
236        };
237        if *step == 0 {
238            return Err(Error::validation(
239                op,
240                ErrorPhase::GraphBuild,
241                ValidationError::InvalidSliceStep { step: 0 },
242            ));
243        }
244        let extent = shape[*axis];
245        if range.start > range.end || range.end > extent {
246            let start = isize::try_from(range.start).map_err(|_| {
247                Error::validation(op, ErrorPhase::GraphBuild, ValidationError::IntegerOverflow)
248            })?;
249            let end = isize::try_from(range.end).map_err(|_| {
250                Error::validation(op, ErrorPhase::GraphBuild, ValidationError::IntegerOverflow)
251            })?;
252            return Err(Error::validation(
253                op,
254                ErrorPhase::GraphBuild,
255                ValidationError::InvalidSliceBounds {
256                    start,
257                    end,
258                    axis_len: extent,
259                },
260            ));
261        }
262        starts[*axis] = range.start;
263        limits[*axis] = range.end;
264        strides[*axis] = *step;
265        has_slice = true;
266    }
267    Ok(has_slice.then_some(SliceConfig {
268        starts,
269        limits,
270        strides,
271    }))
272}
273
274/// Rank-preserving traced tensor slicing builder.
275///
276/// Unspecified axes are kept whole. Range selections become one `Slice`
277/// operation; host-known position selections become `Gather`/`index_select`
278/// operations.
279///
280/// # Examples
281///
282/// ```rust
283/// use tenferro_runtime::{TraceContext, TracedTensor};
284///
285/// let trace = TraceContext::new();
286/// let x = trace.from_vec_col_major(vec![3, 4], vec![0.0_f64; 12]).unwrap();
287/// let y = x.slice_builder().axis(0, 0..2).axis_step(1, 0..4, 2).apply().unwrap();
288/// assert_eq!(y.try_concrete_shape(), Some(vec![2, 2]));
289/// ```
290#[derive(Clone, Debug)]
291pub struct TracedSliceBuilder<'a> {
292    tensor: &'a TracedTensor,
293    selections: Vec<AxisSelection>,
294}
295
296impl<'a> TracedSliceBuilder<'a> {
297    fn new(tensor: &'a TracedTensor) -> Self {
298        Self {
299            tensor,
300            selections: Vec::new(),
301        }
302    }
303
304    /// Add an exclusive-end range selection for one axis.
305    ///
306    /// # Examples
307    ///
308    /// ```rust
309    /// use tenferro_runtime::{TracedTensor};
310    ///
311    /// use tenferro_runtime::TraceContext;
312    /// let trace = TraceContext::new();
313    /// let x = trace.from_vec_col_major(vec![4], vec![1.0_f64, 2.0, 3.0, 4.0]).unwrap();
314    /// let y = x.slice_builder().axis(0, 1..3).apply().unwrap();
315    /// assert_eq!(y.try_concrete_shape(), Some(vec![2]));
316    /// ```
317    pub fn axis(mut self, axis: usize, range: Range<usize>) -> Self {
318        self.selections.push(AxisSelection::Slice {
319            axis,
320            range,
321            step: 1,
322        });
323        self
324    }
325
326    /// Add an exclusive-end strided range selection for one axis.
327    ///
328    /// # Examples
329    ///
330    /// ```rust
331    /// use tenferro_runtime::{TracedTensor};
332    ///
333    /// use tenferro_runtime::TraceContext;
334    /// let trace = TraceContext::new();
335    /// let x = trace.from_vec_col_major(vec![5], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0]).unwrap();
336    /// let y = x.slice_builder().axis_step(0, 0..5, 2).apply().unwrap();
337    /// assert_eq!(y.try_concrete_shape(), Some(vec![3]));
338    /// ```
339    pub fn axis_step(mut self, axis: usize, range: Range<usize>, step: usize) -> Self {
340        self.selections
341            .push(AxisSelection::Slice { axis, range, step });
342        self
343    }
344
345    /// Add a host-known position selection for one axis.
346    ///
347    /// # Examples
348    ///
349    /// ```rust
350    /// use tenferro_runtime::{TracedTensor};
351    ///
352    /// use tenferro_runtime::TraceContext;
353    /// let trace = TraceContext::new();
354    /// let x = trace.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    /// use tenferro_runtime::TraceContext;
374    /// let trace = TraceContext::new();
375    /// let x = trace.from_vec_col_major(vec![4], vec![1.0_f64, 2.0, 3.0, 4.0]).unwrap();
376    /// let y = x.slice_builder().axis(0, 1..4).apply().unwrap();
377    /// assert_eq!(y.try_concrete_shape(), Some(vec![3]));
378    /// ```
379    ///
380    /// # Errors
381    ///
382    /// Returns [`Error::Validation`] with `AxisOutOfBounds` or
383    /// `DuplicateAxis` when selections are invalid, and propagates the
384    /// underlying [`Error::Validation`] from a slice/take graph operation
385    /// that cannot be built.
386    pub fn apply(self) -> Result<TracedTensor> {
387        let shape = concrete_shape_for_axis_slice(self.tensor, "slice_builder")?;
388        let mut seen = vec![false; shape.len()];
389        for selection in &self.selections {
390            let axis = match selection {
391                AxisSelection::Slice { axis, .. } | AxisSelection::Take { axis, .. } => *axis,
392            };
393            validate_axis_selection("slice_builder", shape.len(), &mut seen, axis)?;
394        }
395
396        let mut output = self.tensor.clone();
397        if let Some(config) = apply_slice_axis_config("slice_builder", &shape, &self.selections)? {
398            output = output.slice(config)?;
399        }
400        for selection in self.selections {
401            if let AxisSelection::Take { axis, indices } = selection {
402                output = output.take_axis(axis, &indices)?;
403            }
404        }
405        Ok(output)
406    }
407}
408
409impl TracedTensor {
410    /// Slice one axis with an exclusive-end range, keeping all other axes.
411    ///
412    /// # Examples
413    ///
414    /// ```rust
415    /// use tenferro_runtime::{TracedTensor};
416    ///
417    /// use tenferro_runtime::TraceContext;
418    /// let trace = TraceContext::new();
419    /// let x = trace.from_vec_col_major(vec![4], vec![1.0_f64, 2.0, 3.0, 4.0]).unwrap();
420    /// let y = x.slice_axis(0, 1..3).unwrap();
421    /// assert_eq!(y.try_concrete_shape(), Some(vec![2]));
422    /// ```
423    ///
424    /// # Errors
425    ///
426    /// Returns [`Error::Validation`] with `AxisOutOfBounds` when `axis` is
427    /// outside the concrete rank, or `InvalidArgument` when `range` is
428    /// outside the selected axis extent.
429    pub fn slice_axis(&self, axis: usize, range: Range<usize>) -> Result<Self> {
430        self.slice_builder().axis(axis, range).apply()
431    }
432
433    /// Start a rank-preserving slicing builder for this tensor.
434    ///
435    /// # Examples
436    ///
437    /// ```rust
438    /// use tenferro_runtime::{TracedTensor};
439    ///
440    /// use tenferro_runtime::TraceContext;
441    /// let trace = TraceContext::new();
442    /// let x = trace.from_vec_col_major(vec![3], vec![1.0_f64, 2.0, 3.0]).unwrap();
443    /// let y = x.slice_builder().axis(0, 0..2).apply().unwrap();
444    /// assert_eq!(y.try_concrete_shape(), Some(vec![2]));
445    /// ```
446    pub fn slice_builder(&self) -> TracedSliceBuilder<'_> {
447        TracedSliceBuilder::new(self)
448    }
449
450    /// Select entries from one axis using host-known indices.
451    ///
452    /// # Examples
453    ///
454    /// ```rust
455    /// use tenferro_runtime::{TracedTensor};
456    ///
457    /// use tenferro_runtime::TraceContext;
458    /// let trace = TraceContext::new();
459    /// let x = trace.from_vec_col_major(vec![3], vec![10.0_f64, 20.0, 30.0]).unwrap();
460    /// let y = x.take_axis(0, &[2, 0]).unwrap();
461    /// assert_eq!(y.try_concrete_shape(), Some(vec![2]));
462    /// ```
463    ///
464    /// # Errors
465    ///
466    /// Returns [`Error::Validation`] with `AxisOutOfBounds` when `axis` is
467    /// outside the concrete rank, or `InvalidArgument` when an index list
468    /// cannot be applied to the selected axis.
469    pub fn take_axis(&self, axis: usize, indices: &[usize]) -> Result<Self> {
470        let axis = isize::try_from(axis).map_err(|_| {
471            Error::validation(
472                "take_axis",
473                ErrorPhase::GraphBuild,
474                ValidationError::InvalidArgument {
475                    argument: "axis",
476                    message: format!("axis {axis} cannot be represented as isize"),
477                },
478            )
479        })?;
480        self.index_select(axis, indices)
481    }
482
483    /// Select entries from one axis using host-known positions.
484    ///
485    /// # Examples
486    ///
487    /// ```
488    /// use tenferro_cpu::CpuEngine;
489    /// use tenferro_runtime::{ExecutionEngine, GraphCompiler, Runtime, Tensor, TracedTensor};
490    ///
491    /// use tenferro_runtime::TraceContext;
492    /// let trace = TraceContext::new();
493    /// let x = trace.from_tensor_concrete_shape(
494    ///     Tensor::from_vec_col_major(vec![3], vec![10.0_f64, 20.0, 30.0]).unwrap(),
495    /// )
496    /// .unwrap();
497    /// let y = x.index_select(-1, &[2, 0]).unwrap();
498    /// let mut compiler = GraphCompiler::new();
499    /// let program = compiler.compile(&y).unwrap();
500    /// let mut engine = CpuEngine::serial();
501    /// let mut builder = Runtime::builder();
502    /// builder
503    ///     .register_engine(tenferro_cpu::runtime_engine_registration(&engine).unwrap())
504    ///     .unwrap();
505    /// let runtime = builder.build().unwrap();
506    /// let prepared = runtime
507    ///     .prepare(&program, &[], &[engine.target()])
508    ///     .unwrap();
509    /// let mut outputs = runtime
510    ///     .run(&mut [&mut engine], &prepared, &[])
511    ///     .unwrap();
512    /// let out = outputs.remove(0).into_tensor(&mut engine).unwrap();
513    ///
514    /// assert_eq!(
515    ///     out.as_slice::<f64>().unwrap(),
516    ///     &[30.0, 10.0],
517    /// );
518    /// ```
519    ///
520    /// # Errors
521    ///
522    /// Returns [`Error::Validation`] with `InvalidArgument` when the tensor
523    /// shape is not concrete, `AxisOutOfBounds` when `axis` is outside its
524    /// rank, or `InvalidArgument` when a position is outside the selected
525    /// axis extent.
526    pub fn index_select(&self, axis: isize, positions: &[usize]) -> Result<Self> {
527        let shape = try_concrete_shape(self).ok_or_else(|| {
528            Error::validation(
529                "index_select",
530                ErrorPhase::GraphBuild,
531                ValidationError::InvalidArgument {
532                    argument: "shape",
533                    message: "index_select requires a concrete shape hint".into(),
534                },
535            )
536        })?;
537        let (indices_tensor, config, out_shape) = index_select_config(&shape, axis, positions)?;
538        let indices = self.context().from_tensor_concrete_shape(indices_tensor)?;
539        apply_binary_preserve_input_dtypes(
540            StdTensorOp::Gather(config),
541            self,
542            &indices,
543            out_shape.len(),
544            Some(out_shape.into_iter().map(SymDim::from).collect()),
545            self.dtype,
546        )
547    }
548
549    /// Stack tensors along a newly inserted axis.
550    ///
551    /// # Examples
552    ///
553    /// ```
554    /// use tenferro_cpu::CpuEngine;
555    /// use tenferro_runtime::{ExecutionEngine, GraphCompiler, Runtime, Tensor, TracedTensor};
556    ///
557    /// use tenferro_runtime::TraceContext;
558    /// let trace = TraceContext::new();
559    /// let a = trace.from_tensor_concrete_shape(Tensor::from_vec_col_major(vec![], vec![1.0_f64]).unwrap()).unwrap();
560    /// let b = trace.from_tensor_concrete_shape(Tensor::from_vec_col_major(vec![], vec![2.0_f64]).unwrap()).unwrap();
561    /// let stacked = TracedTensor::stack(&[&a, &b], -1).unwrap();
562    /// let mut compiler = GraphCompiler::new();
563    /// let program = compiler.compile(&stacked).unwrap();
564    /// let mut engine = CpuEngine::serial();
565    /// let mut builder = Runtime::builder();
566    /// builder
567    ///     .register_engine(tenferro_cpu::runtime_engine_registration(&engine).unwrap())
568    ///     .unwrap();
569    /// let runtime = builder.build().unwrap();
570    /// let prepared = runtime
571    ///     .prepare(&program, &[], &[engine.target()])
572    ///     .unwrap();
573    /// let mut outputs = runtime
574    ///     .run(&mut [&mut engine], &prepared, &[])
575    ///     .unwrap();
576    /// let out = outputs.remove(0).into_tensor(&mut engine).unwrap();
577    ///
578    /// assert_eq!(
579    ///     out.as_slice::<f64>().unwrap(),
580    ///     &[1.0, 2.0],
581    /// );
582    /// ```
583    ///
584    /// # Errors
585    ///
586    /// Returns [`Error::Validation`] with `InvalidArgument` for an empty input
587    /// list, `ShapeMismatch` for incompatible input shapes, or
588    /// `AxisOutOfBounds` when `dim` is outside the output rank.
589    pub fn stack(tensors: &[&Self], dim: isize) -> Result<Self> {
590        let first = tensors.first().copied().ok_or_else(|| {
591            Error::validation(
592                "stack",
593                ErrorPhase::GraphBuild,
594                ValidationError::InvalidArgument {
595                    argument: "tensors",
596                    message: "stack requires at least one input".into(),
597                },
598            )
599        })?;
600        let first_shape = try_concrete_shape(first).ok_or_else(|| {
601            Error::validation(
602                "stack",
603                ErrorPhase::GraphBuild,
604                ValidationError::InvalidArgument {
605                    argument: "shape",
606                    message: "stack requires concrete shape hints".into(),
607                },
608            )
609        })?;
610        let mut shapes = Vec::with_capacity(tensors.len());
611        shapes.push(first_shape.as_slice());
612        let mut owned_shapes = Vec::with_capacity(tensors.len().saturating_sub(1));
613        for tensor in tensors.iter().copied().skip(1) {
614            owned_shapes.push(try_concrete_shape(tensor).ok_or_else(|| {
615                Error::validation(
616                    "stack",
617                    ErrorPhase::GraphBuild,
618                    ValidationError::InvalidArgument {
619                        argument: "shape",
620                        message: "stack requires concrete shape hints".into(),
621                    },
622                )
623            })?);
624        }
625        shapes.extend(owned_shapes.iter().map(Vec::as_slice));
626        validate_stack_shapes("stack", &shapes)?;
627
628        let axis = normalize_insert_axis("stack", dim, first.rank)?;
629        let mut expanded_shape = first_shape;
630        expanded_shape.insert(axis, 1);
631        let mut out_shape = expanded_shape.clone();
632        out_shape[axis] = tensors.len();
633        let expanded = tensors
634            .iter()
635            .map(|tensor| tensor.reshape(&expanded_shape))
636            .collect::<Result<Vec<_>>>()?;
637        let refs = expanded.iter().collect::<Vec<_>>();
638        apply_nary_concatenate(
639            &refs,
640            axis,
641            out_shape.into_iter().map(SymDim::from).collect(),
642        )
643    }
644
645    /// Concatenate tensors along one existing axis.
646    ///
647    /// # Errors
648    ///
649    /// Returns [`Error::Validation`] with `InvalidArgument` for an empty input
650    /// list, `RankMismatch`/`ShapeMismatch` for incompatible input shapes, or
651    /// `AxisOutOfBounds` when `axis` is outside the input rank.
652    pub fn concatenate(tensors: &[&Self], axis: usize) -> Result<Self> {
653        let first = tensors.first().copied().ok_or_else(|| {
654            Error::validation(
655                "concatenate",
656                ErrorPhase::GraphBuild,
657                ValidationError::InvalidArgument {
658                    argument: "tensors",
659                    message: "concatenate requires at least one input".into(),
660                },
661            )
662        })?;
663        if axis >= first.rank {
664            return Err(Error::validation(
665                "concatenate",
666                ErrorPhase::GraphBuild,
667                ValidationError::AxisOutOfBounds {
668                    axis,
669                    rank: first.rank,
670                },
671            ));
672        }
673        for tensor in tensors.iter().copied().skip(1) {
674            if tensor.rank != first.rank {
675                return Err(Error::validation(
676                    "concatenate",
677                    ErrorPhase::GraphBuild,
678                    ValidationError::RankMismatch {
679                        expected: first.rank,
680                        actual: tensor.rank,
681                    },
682                ));
683            }
684        }
685
686        let op = StdTensorOp::Concatenate {
687            axis,
688            input_count: tensors.len(),
689        };
690        let (_, out_shape_hint) =
691            infer_traced_single_output_shape("TracedTensor::concatenate", &op, tensors)?;
692        let out_shape = out_shape_hint.ok_or_else(|| {
693            Error::Internal("concatenate shape inference returned no shape hint".into())
694        })?;
695        apply_nary_concatenate(tensors, axis, out_shape)
696    }
697}
698
699fn apply_nary_concatenate(
700    tensors: &[&TracedTensor],
701    axis: usize,
702    out_shape: Vec<SymDim>,
703) -> Result<TracedTensor> {
704    let out_dtype = promote_dtypes(tensors.iter().map(|tensor| tensor.dtype));
705    let tensors = tensors
706        .iter()
707        .map(|tensor| {
708            if tensor.dtype != out_dtype {
709                tensor.cast(out_dtype)
710            } else {
711                Ok((*tensor).clone())
712            }
713        })
714        .collect::<Result<Vec<_>>>()?;
715
716    let Some(first) = tensors.first() else {
717        return Err(Error::invalid_argument(
718            "TracedSliceBuilder::concatenate",
719            ErrorPhase::GraphBuild,
720            "tensors",
721            "at least one tensor is required",
722        ));
723    };
724    for tensor in tensors.iter().skip(1) {
725        crate::traced::ensure_same_context(first, tensor)?;
726    }
727    let values: Vec<_> = tensors.iter().map(|tensor| tensor.value).collect();
728    let op = StdTensorOp::Concatenate {
729        axis,
730        input_count: tensors.len(),
731    };
732    let outputs = crate::trace_ops::append_core(&first.ctx, &op, &values)?;
733    crate::trace_ops::finish_value(
734        &first.ctx,
735        outputs[0],
736        out_shape.len(),
737        out_dtype,
738        Some(out_shape),
739    )
740}
741
742#[cfg(test)]
743mod tests;
tenferro_runtime/shape_packing.rs

tenferro_runtime/
shape_packing.rs
