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