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::{GatherConfig, SliceConfig, Tensor, TypedTensor};
8
9use crate::checkpoint::CheckpointNode;
10use crate::error::{Error, Result};
11use crate::metadata::{register_scoped_value_metadata, tensor_meta, MetadataScopeChain};
12use crate::shape_infer::promote_dtypes;
13use crate::sym_dim::SymDim;
14use crate::traced::{
15    apply_binary_preserve_input_dtypes, infer_traced_single_output_shape, merge_traced_inputs_map,
16    next_traced_id, try_concrete_shape,
17};
18use crate::TracedTensor;
19
20fn normalize_existing_axis(op: &'static str, axis: isize, rank: usize) -> Result<usize> {
21    let normalized = if axis >= 0 {
22        axis as usize
23    } else {
24        rank.checked_sub(axis.unsigned_abs())
25            .ok_or(tenferro_tensor::Error::AxisOutOfBounds {
26                op,
27                axis: axis.unsigned_abs(),
28                rank,
29            })?
30    };
31    if normalized >= rank {
32        return Err(tenferro_tensor::Error::AxisOutOfBounds {
33            op,
34            axis: axis.unsigned_abs(),
35            rank,
36        }
37        .into());
38    }
39    Ok(normalized)
40}
41
42fn normalize_insert_axis(op: &'static str, axis: isize, rank: usize) -> Result<usize> {
43    let insert_rank = rank
44        .checked_add(1)
45        .ok_or(tenferro_tensor::Error::AxisOutOfBounds {
46            op,
47            axis: axis.unsigned_abs(),
48            rank,
49        })?;
50    let normalized = if axis >= 0 {
51        axis as usize
52    } else {
53        insert_rank.checked_sub(axis.unsigned_abs()).ok_or(
54            tenferro_tensor::Error::AxisOutOfBounds {
55                op,
56                axis: axis.unsigned_abs(),
57                rank: insert_rank,
58            },
59        )?
60    };
61    if normalized > rank {
62        return Err(tenferro_tensor::Error::AxisOutOfBounds {
63            op,
64            axis: axis.unsigned_abs(),
65            rank: insert_rank,
66        }
67        .into());
68    }
69    Ok(normalized)
70}
71
72fn index_select_config(
73    shape: &[usize],
74    axis: isize,
75    positions: &[usize],
76) -> Result<(Tensor, GatherConfig, Vec<usize>)> {
77    let axis = normalize_existing_axis("index_select", axis, shape.len())?;
78    let axis_extent = shape[axis];
79    for &position in positions {
80        if position >= axis_extent {
81            return Err(tenferro_tensor::Error::InvalidConfig {
82                op: "index_select",
83                message: format!(
84                    "position {position} out of bounds for axis {axis} with extent {axis_extent}"
85                ),
86            }
87            .into());
88        }
89    }
90
91    let mut out_shape = shape.to_vec();
92    out_shape[axis] = positions.len();
93
94    let mut slice_sizes = shape.to_vec();
95    slice_sizes[axis] = 1;
96
97    let offset_dims = (0..shape.len()).filter(|&dim| dim != axis).collect();
98    let index_data = positions
99        .iter()
100        .map(|&position| {
101            i64::try_from(position).map_err(|_| tenferro_tensor::Error::InvalidConfig {
102                op: "index_select",
103                message: format!("position {position} cannot be represented as i64"),
104            })
105        })
106        .collect::<tenferro_tensor::Result<Vec<_>>>()?;
107    let indices = Tensor::I64(TypedTensor::from_vec_col_major(
108        vec![positions.len(), 1],
109        index_data,
110    )?);
111
112    let config = GatherConfig {
113        offset_dims,
114        collapsed_slice_dims: vec![axis],
115        start_index_map: vec![axis],
116        index_vector_dim: 1,
117        slice_sizes,
118    };
119
120    Ok((indices, config, out_shape))
121}
122
123fn validate_stack_shapes(op: &'static str, shapes: &[&[usize]]) -> Result<()> {
124    let Some(first) = shapes.first() else {
125        return Err(tenferro_tensor::Error::InvalidConfig {
126            op,
127            message: "stack requires at least one input".into(),
128        }
129        .into());
130    };
131    for shape in shapes.iter().skip(1) {
132        if *shape != *first {
133            return Err(tenferro_tensor::Error::ShapeMismatch {
134                op,
135                lhs: first.to_vec(),
136                rhs: shape.to_vec(),
137            }
138            .into());
139        }
140    }
141    Ok(())
142}
143
144#[derive(Clone, Debug)]
145enum AxisSelection {
146    Slice {
147        axis: usize,
148        range: Range<usize>,
149        step: usize,
150    },
151    Take {
152        axis: usize,
153        indices: Vec<usize>,
154    },
155}
156
157fn concrete_shape_for_axis_slice(tensor: &TracedTensor, op: &'static str) -> Result<Vec<usize>> {
158    try_concrete_shape(tensor).ok_or_else(|| Error::InvalidGraphBuild {
159        op,
160        message: format!("{op} requires a concrete shape hint"),
161    })
162}
163
164fn validate_axis_selection(
165    op: &'static str,
166    rank: usize,
167    seen: &mut [bool],
168    axis: usize,
169) -> Result<()> {
170    if axis >= rank {
171        return Err(tenferro_tensor::Error::AxisOutOfBounds { op, axis, rank }.into());
172    }
173    if seen[axis] {
174        return Err(tenferro_tensor::Error::DuplicateAxis {
175            op,
176            axis,
177            role: "selection",
178        }
179        .into());
180    }
181    seen[axis] = true;
182    Ok(())
183}
184
185fn apply_slice_axis_config(
186    op: &'static str,
187    shape: &[usize],
188    selections: &[AxisSelection],
189) -> Result<Option<SliceConfig>> {
190    let mut starts = vec![0; shape.len()];
191    let mut limits = shape.to_vec();
192    let mut strides = vec![1; shape.len()];
193    let mut has_slice = false;
194    for selection in selections {
195        let AxisSelection::Slice { axis, range, step } = selection else {
196            continue;
197        };
198        if *step == 0 {
199            return Err(tenferro_tensor::Error::InvalidConfig {
200                op,
201                message: format!("axis {axis} has zero step"),
202            }
203            .into());
204        }
205        let extent = shape[*axis];
206        if range.start > range.end || range.end > extent {
207            return Err(tenferro_tensor::Error::InvalidConfig {
208                op,
209                message: format!(
210                    "axis {axis} range {}..{} is out of bounds for extent {extent}",
211                    range.start, range.end
212                ),
213            }
214            .into());
215        }
216        starts[*axis] = range.start;
217        limits[*axis] = range.end;
218        strides[*axis] = *step;
219        has_slice = true;
220    }
221    Ok(has_slice.then_some(SliceConfig {
222        starts,
223        limits,
224        strides,
225    }))
226}
227
228/// Rank-preserving traced tensor slicing builder.
229///
230/// Unspecified axes are kept whole. Range selections become one `Slice`
231/// operation; host-known position selections become `Gather`/`index_select`
232/// operations.
233///
234/// # Examples
235///
236/// ```rust
237/// use tenferro_runtime::TracedTensor;
238///
239/// let x = TracedTensor::from_vec_col_major(vec![3, 4], vec![0.0_f64; 12]).unwrap();
240/// let y = x.slice_builder().axis(0, 0..2).axis_step(1, 0..4, 2).apply().unwrap();
241/// assert_eq!(y.try_concrete_shape(), Some(vec![2, 2]));
242/// ```
243#[derive(Clone, Debug)]
244pub struct TracedSliceBuilder<'a> {
245    tensor: &'a TracedTensor,
246    selections: Vec<AxisSelection>,
247}
248
249impl<'a> TracedSliceBuilder<'a> {
250    fn new(tensor: &'a TracedTensor) -> Self {
251        Self {
252            tensor,
253            selections: Vec::new(),
254        }
255    }
256
257    /// Add an exclusive-end range selection for one axis.
258    ///
259    /// # Examples
260    ///
261    /// ```rust
262    /// use tenferro_runtime::TracedTensor;
263    ///
264    /// let x = TracedTensor::from_vec_col_major(vec![4], vec![1.0_f64, 2.0, 3.0, 4.0]).unwrap();
265    /// let y = x.slice_builder().axis(0, 1..3).apply().unwrap();
266    /// assert_eq!(y.try_concrete_shape(), Some(vec![2]));
267    /// ```
268    pub fn axis(mut self, axis: usize, range: Range<usize>) -> Self {
269        self.selections.push(AxisSelection::Slice {
270            axis,
271            range,
272            step: 1,
273        });
274        self
275    }
276
277    /// Add an exclusive-end strided range selection for one axis.
278    ///
279    /// # Examples
280    ///
281    /// ```rust
282    /// use tenferro_runtime::TracedTensor;
283    ///
284    /// let x = TracedTensor::from_vec_col_major(vec![5], vec![1.0_f64, 2.0, 3.0, 4.0, 5.0]).unwrap();
285    /// let y = x.slice_builder().axis_step(0, 0..5, 2).apply().unwrap();
286    /// assert_eq!(y.try_concrete_shape(), Some(vec![3]));
287    /// ```
288    pub fn axis_step(mut self, axis: usize, range: Range<usize>, step: usize) -> Self {
289        self.selections
290            .push(AxisSelection::Slice { axis, range, step });
291        self
292    }
293
294    /// Add a host-known position selection for one axis.
295    ///
296    /// # Examples
297    ///
298    /// ```rust
299    /// use tenferro_runtime::TracedTensor;
300    ///
301    /// let x = TracedTensor::from_vec_col_major(vec![3], vec![1.0_f64, 2.0, 3.0]).unwrap();
302    /// let y = x.slice_builder().take_axis(0, &[2, 0]).apply().unwrap();
303    /// assert_eq!(y.try_concrete_shape(), Some(vec![2]));
304    /// ```
305    pub fn take_axis(mut self, axis: usize, indices: &[usize]) -> Self {
306        self.selections.push(AxisSelection::Take {
307            axis,
308            indices: indices.to_vec(),
309        });
310        self
311    }
312
313    /// Build and apply the requested slice/take operations.
314    ///
315    /// # Examples
316    ///
317    /// ```rust
318    /// use tenferro_runtime::TracedTensor;
319    ///
320    /// let x = TracedTensor::from_vec_col_major(vec![4], vec![1.0_f64, 2.0, 3.0, 4.0]).unwrap();
321    /// let y = x.slice_builder().axis(0, 1..4).apply().unwrap();
322    /// assert_eq!(y.try_concrete_shape(), Some(vec![3]));
323    /// ```
324    pub fn apply(self) -> Result<TracedTensor> {
325        let shape = concrete_shape_for_axis_slice(self.tensor, "slice_builder")?;
326        let mut seen = vec![false; shape.len()];
327        for selection in &self.selections {
328            let axis = match selection {
329                AxisSelection::Slice { axis, .. } | AxisSelection::Take { axis, .. } => *axis,
330            };
331            validate_axis_selection("slice_builder", shape.len(), &mut seen, axis)?;
332        }
333
334        let mut output = self.tensor.clone();
335        if let Some(config) = apply_slice_axis_config("slice_builder", &shape, &self.selections)? {
336            output = output.slice(config)?;
337        }
338        for selection in self.selections {
339            if let AxisSelection::Take { axis, indices } = selection {
340                output = output.take_axis(axis, &indices)?;
341            }
342        }
343        Ok(output)
344    }
345}
346
347impl TracedTensor {
348    /// Slice one axis with an exclusive-end range, keeping all other axes.
349    ///
350    /// # Examples
351    ///
352    /// ```rust
353    /// use tenferro_runtime::TracedTensor;
354    ///
355    /// let x = TracedTensor::from_vec_col_major(vec![4], vec![1.0_f64, 2.0, 3.0, 4.0]).unwrap();
356    /// let y = x.slice_axis(0, 1..3).unwrap();
357    /// assert_eq!(y.try_concrete_shape(), Some(vec![2]));
358    /// ```
359    pub fn slice_axis(&self, axis: usize, range: Range<usize>) -> Result<Self> {
360        self.slice_builder().axis(axis, range).apply()
361    }
362
363    /// Start a rank-preserving slicing builder for this tensor.
364    ///
365    /// # Examples
366    ///
367    /// ```rust
368    /// use tenferro_runtime::TracedTensor;
369    ///
370    /// let x = TracedTensor::from_vec_col_major(vec![3], vec![1.0_f64, 2.0, 3.0]).unwrap();
371    /// let y = x.slice_builder().axis(0, 0..2).apply().unwrap();
372    /// assert_eq!(y.try_concrete_shape(), Some(vec![2]));
373    /// ```
374    pub fn slice_builder(&self) -> TracedSliceBuilder<'_> {
375        TracedSliceBuilder::new(self)
376    }
377
378    /// Select entries from one axis using host-known indices.
379    ///
380    /// # Examples
381    ///
382    /// ```rust
383    /// use tenferro_runtime::TracedTensor;
384    ///
385    /// let x = TracedTensor::from_vec_col_major(vec![3], vec![10.0_f64, 20.0, 30.0]).unwrap();
386    /// let y = x.take_axis(0, &[2, 0]).unwrap();
387    /// assert_eq!(y.try_concrete_shape(), Some(vec![2]));
388    /// ```
389    pub fn take_axis(&self, axis: usize, indices: &[usize]) -> Result<Self> {
390        let axis = isize::try_from(axis).map_err(|_| {
391            Error::TensorRuntime(tenferro_tensor::Error::InvalidConfig {
392                op: "take_axis",
393                message: format!("axis {axis} cannot be represented as isize"),
394            })
395        })?;
396        self.index_select(axis, indices)
397    }
398
399    /// Select entries from one axis using host-known positions.
400    ///
401    /// # Examples
402    ///
403    /// ```
404    /// use tenferro_cpu::CpuBackend;
405    /// use tenferro_runtime::{GraphCompiler, GraphExecutor, Tensor, TracedTensor};
406    ///
407    /// let x = TracedTensor::from_tensor_concrete_shape(
408    ///     Tensor::from_vec_col_major(vec![3], vec![10.0_f64, 20.0, 30.0]).unwrap(),
409    /// )
410    /// .unwrap();
411    /// let y = x.index_select(-1, &[2, 0]).unwrap();
412    /// let mut compiler = GraphCompiler::new();
413    /// let program = compiler.compile(&y).unwrap();
414    /// let out = GraphExecutor::new(CpuBackend::new()).run(&program).unwrap();
415    ///
416    /// assert_eq!(
417    ///     out.as_slice::<f64>().unwrap(),
418    ///     &[30.0, 10.0],
419    /// );
420    /// ```
421    pub fn index_select(&self, axis: isize, positions: &[usize]) -> Result<Self> {
422        let shape = try_concrete_shape(self).ok_or_else(|| Error::InvalidGraphBuild {
423            op: "index_select",
424            message: "index_select requires a concrete shape hint".into(),
425        })?;
426        let (indices_tensor, config, out_shape) = index_select_config(&shape, axis, positions)?;
427        let indices = TracedTensor::from_tensor_concrete_shape(indices_tensor)?;
428        apply_binary_preserve_input_dtypes(
429            StdTensorOp::Gather(config),
430            self,
431            &indices,
432            out_shape.len(),
433            Some(out_shape.into_iter().map(SymDim::from).collect()),
434            self.dtype,
435        )
436    }
437
438    /// Stack tensors along a newly inserted axis.
439    ///
440    /// # Examples
441    ///
442    /// ```
443    /// use tenferro_cpu::CpuBackend;
444    /// use tenferro_runtime::{GraphCompiler, GraphExecutor, Tensor, TracedTensor};
445    ///
446    /// let a = TracedTensor::from_tensor_concrete_shape(Tensor::from_vec_col_major(vec![], vec![1.0_f64]).unwrap()).unwrap();
447    /// let b = TracedTensor::from_tensor_concrete_shape(Tensor::from_vec_col_major(vec![], vec![2.0_f64]).unwrap()).unwrap();
448    /// let stacked = TracedTensor::stack(&[&a, &b], -1).unwrap();
449    /// let mut compiler = GraphCompiler::new();
450    /// let program = compiler.compile(&stacked).unwrap();
451    /// let out = GraphExecutor::new(CpuBackend::new()).run(&program).unwrap();
452    ///
453    /// assert_eq!(
454    ///     out.as_slice::<f64>().unwrap(),
455    ///     &[1.0, 2.0],
456    /// );
457    /// ```
458    pub fn stack(tensors: &[&Self], dim: isize) -> Result<Self> {
459        let first = tensors.first().copied().ok_or_else(|| {
460            Error::TensorRuntime(tenferro_tensor::Error::InvalidConfig {
461                op: "stack",
462                message: "stack requires at least one input".into(),
463            })
464        })?;
465        let first_shape = try_concrete_shape(first).ok_or_else(|| Error::InvalidGraphBuild {
466            op: "stack",
467            message: "stack requires concrete shape hints".into(),
468        })?;
469        let mut shapes = Vec::with_capacity(tensors.len());
470        shapes.push(first_shape.as_slice());
471        let mut owned_shapes = Vec::with_capacity(tensors.len().saturating_sub(1));
472        for tensor in tensors.iter().copied().skip(1) {
473            owned_shapes.push(try_concrete_shape(tensor).ok_or_else(|| {
474                Error::InvalidGraphBuild {
475                    op: "stack",
476                    message: "stack requires concrete shape hints".into(),
477                }
478            })?);
479        }
480        shapes.extend(owned_shapes.iter().map(Vec::as_slice));
481        validate_stack_shapes("stack", &shapes)?;
482
483        let axis = normalize_insert_axis("stack", dim, first.rank)?;
484        let mut expanded_shape = first_shape;
485        expanded_shape.insert(axis, 1);
486        let mut out_shape = expanded_shape.clone();
487        out_shape[axis] = tensors.len();
488        let expanded = tensors
489            .iter()
490            .map(|tensor| tensor.reshape(&expanded_shape))
491            .collect::<Result<Vec<_>>>()?;
492        let refs = expanded.iter().collect::<Vec<_>>();
493        apply_nary_concatenate(
494            &refs,
495            axis,
496            out_shape.into_iter().map(SymDim::from).collect(),
497        )
498    }
499
500    /// Concatenate tensors along one existing axis.
501    pub fn concatenate(tensors: &[&Self], axis: usize) -> Result<Self> {
502        let first = tensors.first().copied().ok_or_else(|| {
503            Error::TensorRuntime(tenferro_tensor::Error::InvalidConfig {
504                op: "concatenate",
505                message: "concatenate requires at least one input".into(),
506            })
507        })?;
508        if axis >= first.rank {
509            return Err(tenferro_tensor::Error::AxisOutOfBounds {
510                op: "concatenate",
511                axis,
512                rank: first.rank,
513            }
514            .into());
515        }
516        for tensor in tensors.iter().copied().skip(1) {
517            if tensor.rank != first.rank {
518                return Err(tenferro_tensor::Error::RankMismatch {
519                    op: "concatenate",
520                    expected: first.rank,
521                    actual: tensor.rank,
522                }
523                .into());
524            }
525        }
526
527        let op = StdTensorOp::Concatenate {
528            axis,
529            input_count: tensors.len(),
530        };
531        let (_, out_shape_hint) =
532            infer_traced_single_output_shape("TracedTensor::concatenate", &op, tensors)?;
533        let out_shape = out_shape_hint.ok_or_else(|| {
534            Error::Internal("concatenate shape inference returned no shape hint".into())
535        })?;
536        apply_nary_concatenate(tensors, axis, out_shape)
537    }
538}
539
540fn apply_nary_concatenate(
541    tensors: &[&TracedTensor],
542    axis: usize,
543    out_shape: Vec<SymDim>,
544) -> Result<TracedTensor> {
545    let out_dtype = promote_dtypes(tensors.iter().map(|tensor| tensor.dtype));
546    let tensors = tensors
547        .iter()
548        .map(|tensor| {
549            if tensor.dtype != out_dtype {
550                tensor.cast(out_dtype)
551            } else {
552                Ok((*tensor).clone())
553            }
554        })
555        .collect::<Result<Vec<_>>>()?;
556
557    let mut builder = GraphBuilder::new();
558    for tensor in &tensors {
559        builder.add_parent(tensor.graph.clone());
560    }
561    let input_refs = tensors
562        .iter()
563        .map(|tensor| ValueRef::External(tensor.graph.values()[tensor.val].key.clone()))
564        .collect::<Vec<_>>();
565    let outputs = builder.add_operation(
566        StdTensorOp::Concatenate {
567            axis,
568            input_count: tensors.len(),
569        },
570        input_refs,
571        OperationRole::Primary,
572    );
573    builder.set_outputs(outputs.clone());
574    let graph = Arc::new(builder.build());
575    // Callers route through shape inference before graph construction.
576    let metadata_scope =
577        super::traced::register_metadata_or_internal(register_scoped_value_metadata(
578            graph.values()[outputs[0]].key.clone(),
579            tensor_meta(out_dtype, out_shape.clone()),
580        ))?;
581
582    let inputs_map = merge_traced_inputs_map(tensors.iter());
583    let mut extra_roots = Vec::new();
584    let mut checkpoint_chain = None;
585    for tensor in &tensors {
586        extra_roots.extend(tensor.extra_roots.iter().cloned());
587        checkpoint_chain =
588            CheckpointNode::merge_chains(checkpoint_chain, tensor.checkpoint_chain.clone());
589    }
590    Ok(TracedTensor {
591        id: next_traced_id(),
592        rank: out_shape.len(),
593        dtype: out_dtype,
594        graph,
595        val: outputs[0],
596        data: None,
597        shape_hint: Some(out_shape),
598        inputs_map,
599        extra_roots,
600        checkpoint_chain,
601        metadata_scopes: MetadataScopeChain::with_new(
602            metadata_scope,
603            tensors.iter().map(|tensor| &tensor.metadata_scopes),
604        ),
605    })
606}
607
608#[cfg(test)]
609mod tests {
610    use super::{normalize_existing_axis, normalize_insert_axis};
611
612    #[test]
613    fn axis_normalization_handles_ranks_larger_than_isize_max() {
614        assert_eq!(normalize_existing_axis("test", 0, usize::MAX).unwrap(), 0);
615        assert_eq!(
616            normalize_existing_axis("test", -1, usize::MAX).unwrap(),
617            usize::MAX - 1
618        );
619        assert_eq!(
620            normalize_insert_axis("test", -1, usize::MAX - 1).unwrap(),
621            usize::MAX - 1
622        );
623        assert!(normalize_insert_axis("test", -1, usize::MAX).is_err());
624    }
625}
tenferro_runtime/shape_packing.rs

tenferro_runtime/
shape_packing.rs
