tenferro_runtime/
shape_packing.rs

1use std::collections::HashMap;
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, Tensor, TypedTensor};
8
9use crate::checkpoint::CheckpointNode;
10use crate::error::{Error, Result};
11use crate::metadata::{metadata_scopes_with_new, register_scoped_value_metadata, tensor_meta};
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, next_traced_id,
16    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
144impl TracedTensor {
145    /// Select entries from one axis using host-known positions.
146    ///
147    /// # Examples
148    ///
149    /// ```
150    /// use tenferro_cpu::CpuBackend;
151    /// use tenferro_runtime::{GraphCompiler, GraphExecutor, Tensor, TracedTensor};
152    ///
153    /// let x = TracedTensor::from_tensor_concrete_shape(
154    ///     Tensor::from_vec_col_major(vec![3], vec![10.0_f64, 20.0, 30.0]).unwrap(),
155    /// )
156    /// .unwrap();
157    /// let y = x.index_select(-1, &[2, 0]).unwrap();
158    /// let mut compiler = GraphCompiler::new();
159    /// let program = compiler.compile(&y).unwrap();
160    /// let out = GraphExecutor::new(CpuBackend::new()).run(&program).unwrap();
161    ///
162    /// assert_eq!(
163    ///     out.as_slice::<f64>().unwrap(),
164    ///     &[30.0, 10.0],
165    /// );
166    /// ```
167    pub fn index_select(&self, axis: isize, positions: &[usize]) -> Result<Self> {
168        let shape = try_concrete_shape(self).ok_or_else(|| Error::InvalidGraphBuild {
169            op: "index_select",
170            message: "index_select requires a concrete shape hint".into(),
171        })?;
172        let (indices_tensor, config, out_shape) = index_select_config(&shape, axis, positions)?;
173        let indices = TracedTensor::from_tensor_concrete_shape(indices_tensor)?;
174        Ok(apply_binary_preserve_input_dtypes(
175            StdTensorOp::Gather(config),
176            self,
177            &indices,
178            out_shape.len(),
179            Some(out_shape.into_iter().map(SymDim::from).collect()),
180            self.dtype,
181        ))
182    }
183
184    /// Stack tensors along a newly inserted axis.
185    ///
186    /// # Examples
187    ///
188    /// ```
189    /// use tenferro_cpu::CpuBackend;
190    /// use tenferro_runtime::{GraphCompiler, GraphExecutor, Tensor, TracedTensor};
191    ///
192    /// let a = TracedTensor::from_tensor_concrete_shape(Tensor::from_vec_col_major(vec![], vec![1.0_f64]).unwrap()).unwrap();
193    /// let b = TracedTensor::from_tensor_concrete_shape(Tensor::from_vec_col_major(vec![], vec![2.0_f64]).unwrap()).unwrap();
194    /// let stacked = TracedTensor::stack(&[&a, &b], -1).unwrap();
195    /// let mut compiler = GraphCompiler::new();
196    /// let program = compiler.compile(&stacked).unwrap();
197    /// let out = GraphExecutor::new(CpuBackend::new()).run(&program).unwrap();
198    ///
199    /// assert_eq!(
200    ///     out.as_slice::<f64>().unwrap(),
201    ///     &[1.0, 2.0],
202    /// );
203    /// ```
204    pub fn stack(tensors: &[&Self], dim: isize) -> Result<Self> {
205        let first = tensors.first().copied().ok_or_else(|| {
206            Error::TensorRuntime(tenferro_tensor::Error::InvalidConfig {
207                op: "stack",
208                message: "stack requires at least one input".into(),
209            })
210        })?;
211        let first_shape = try_concrete_shape(first).ok_or_else(|| Error::InvalidGraphBuild {
212            op: "stack",
213            message: "stack requires concrete shape hints".into(),
214        })?;
215        let mut shapes = Vec::with_capacity(tensors.len());
216        shapes.push(first_shape.as_slice());
217        let mut owned_shapes = Vec::with_capacity(tensors.len().saturating_sub(1));
218        for tensor in tensors.iter().copied().skip(1) {
219            owned_shapes.push(try_concrete_shape(tensor).ok_or_else(|| {
220                Error::InvalidGraphBuild {
221                    op: "stack",
222                    message: "stack requires concrete shape hints".into(),
223                }
224            })?);
225        }
226        shapes.extend(owned_shapes.iter().map(Vec::as_slice));
227        validate_stack_shapes("stack", &shapes)?;
228
229        let axis = normalize_insert_axis("stack", dim, first.rank)?;
230        let mut expanded_shape = first_shape;
231        expanded_shape.insert(axis, 1);
232        let mut out_shape = expanded_shape.clone();
233        out_shape[axis] = tensors.len();
234        let expanded = tensors
235            .iter()
236            .map(|tensor| tensor.reshape(&expanded_shape))
237            .collect::<Vec<_>>();
238        let refs = expanded.iter().collect::<Vec<_>>();
239        Ok(apply_nary_concatenate(
240            &refs,
241            axis,
242            out_shape.into_iter().map(SymDim::from).collect(),
243        ))
244    }
245
246    /// Concatenate tensors along one existing axis.
247    pub fn concatenate(tensors: &[&Self], axis: usize) -> Result<Self> {
248        let first = tensors.first().copied().ok_or_else(|| {
249            Error::TensorRuntime(tenferro_tensor::Error::InvalidConfig {
250                op: "concatenate",
251                message: "concatenate requires at least one input".into(),
252            })
253        })?;
254        if axis >= first.rank {
255            return Err(tenferro_tensor::Error::AxisOutOfBounds {
256                op: "concatenate",
257                axis,
258                rank: first.rank,
259            }
260            .into());
261        }
262        for tensor in tensors.iter().copied().skip(1) {
263            if tensor.rank != first.rank {
264                return Err(tenferro_tensor::Error::RankMismatch {
265                    op: "concatenate",
266                    expected: first.rank,
267                    actual: tensor.rank,
268                }
269                .into());
270            }
271        }
272
273        let op = StdTensorOp::Concatenate {
274            axis,
275            input_count: tensors.len(),
276        };
277        let (_, out_shape_hint) =
278            infer_traced_single_output_shape("TracedTensor::concatenate", &op, tensors)?;
279        let out_shape = out_shape_hint.ok_or_else(|| {
280            Error::Internal("concatenate shape inference returned no shape hint".into())
281        })?;
282        Ok(apply_nary_concatenate(tensors, axis, out_shape))
283    }
284}
285
286fn apply_nary_concatenate(
287    tensors: &[&TracedTensor],
288    axis: usize,
289    out_shape: Vec<SymDim>,
290) -> TracedTensor {
291    let out_dtype = promote_dtypes(tensors.iter().map(|tensor| tensor.dtype));
292    let tensors = tensors
293        .iter()
294        .map(|tensor| {
295            if tensor.dtype != out_dtype {
296                tensor.cast(out_dtype)
297            } else {
298                (*tensor).clone()
299            }
300        })
301        .collect::<Vec<_>>();
302
303    let mut builder = GraphBuilder::new();
304    for tensor in &tensors {
305        builder.add_parent(tensor.graph.clone());
306    }
307    let input_refs = tensors
308        .iter()
309        .map(|tensor| ValueRef::External(tensor.graph.values()[tensor.val].key.clone()))
310        .collect::<Vec<_>>();
311    let outputs = builder.add_operation(
312        StdTensorOp::Concatenate {
313            axis,
314            input_count: tensors.len(),
315        },
316        input_refs,
317        OperationRole::Primary,
318    );
319    builder.set_outputs(outputs.clone());
320    let graph = Arc::new(builder.build());
321    // Callers route through shape inference before graph construction.
322    let metadata_scope = register_scoped_value_metadata(
323        graph.values()[outputs[0]].key.clone(),
324        tensor_meta(out_dtype, out_shape.clone()),
325    )
326    .expect("fresh concatenate output metadata registration failed");
327
328    let mut inputs_map = HashMap::new();
329    let mut extra_roots = Vec::new();
330    let mut checkpoint_chain = None;
331    for tensor in &tensors {
332        inputs_map.extend(
333            tensor
334                .inputs_map
335                .iter()
336                .map(|(k, v)| (k.clone(), v.clone())),
337        );
338        extra_roots.extend(tensor.extra_roots.iter().cloned());
339        checkpoint_chain =
340            CheckpointNode::merge_chains(checkpoint_chain, tensor.checkpoint_chain.clone());
341    }
342    let inherited_scopes = tensors
343        .iter()
344        .map(|tensor| tensor.metadata_scopes.as_slice())
345        .collect::<Vec<_>>();
346
347    TracedTensor {
348        id: next_traced_id(),
349        rank: out_shape.len(),
350        dtype: out_dtype,
351        graph,
352        val: outputs[0],
353        data: None,
354        shape_hint: Some(out_shape),
355        inputs_map: Arc::new(inputs_map),
356        extra_roots,
357        checkpoint_chain,
358        metadata_scopes: metadata_scopes_with_new(metadata_scope, inherited_scopes),
359    }
360}
361
362#[cfg(test)]
363mod tests {
364    use super::{normalize_existing_axis, normalize_insert_axis};
365
366    #[test]
367    fn axis_normalization_handles_ranks_larger_than_isize_max() {
368        assert_eq!(normalize_existing_axis("test", 0, usize::MAX).unwrap(), 0);
369        assert_eq!(
370            normalize_existing_axis("test", -1, usize::MAX).unwrap(),
371            usize::MAX - 1
372        );
373        assert_eq!(
374            normalize_insert_axis("test", -1, usize::MAX - 1).unwrap(),
375            usize::MAX - 1
376        );
377        assert!(normalize_insert_axis("test", -1, usize::MAX).is_err());
378    }
379}
tenferro_runtime/shape_packing.rs

tenferro_runtime/
shape_packing.rs
