tenferro_ad/
shape_packing.rs

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

tenferro_ad/
shape_packing.rs
