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

tenferro_ad/
shape_packing.rs
