rlx_tensor/
slice.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3
4//! Slice specs — view-like narrowing without materializing data.
5
6use rlx_ir::Dim;
7
8use crate::Tensor;
9
10/// One axis in a slice (`s![ax(), rg(2, 10)]` or `s![.., 2..10]` via helpers).
11///
12/// Indices are signed and NumPy-style: negative values count from the end
13/// (`-1` = last), resolved against the axis's (static) size.
14#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15pub enum SliceAxis {
16    All,
17    Index(i64),
18    Range { start: i64, end: i64 },
19    From(i64),
20}
21
22/// Keep the full axis.
23#[inline]
24pub fn ax() -> SliceAxis {
25    SliceAxis::All
26}
27
28/// Single index (negative counts from the end). Keeps the axis as size 1 —
29/// use [`Tensor::select`](crate::Tensor::select) to drop it.
30#[inline]
31pub fn ix(i: i64) -> SliceAxis {
32    SliceAxis::Index(i)
33}
34
35/// Half-open range `start..end` (either bound may be negative).
36#[inline]
37pub fn rg(start: i64, end: i64) -> SliceAxis {
38    SliceAxis::Range { start, end }
39}
40
41/// `start..` through the end of a static axis (negative `start` allowed).
42#[inline]
43pub fn tail(start: i64) -> SliceAxis {
44    SliceAxis::From(start)
45}
46
47/// Resolve a signed index against an axis size (negative → `size + idx`).
48/// Panics on a dynamic axis or out-of-range index.
49pub(crate) fn resolve_index(idx: i64, dim: Option<usize>, axis: usize) -> usize {
50    if idx >= 0 {
51        return idx as usize;
52    }
53    let n = dim
54        .unwrap_or_else(|| panic!("negative index requires a static dimension on axis {axis}"))
55        as i64;
56    let resolved = n + idx;
57    assert!(
58        resolved >= 0,
59        "index {idx} out of bounds for axis {axis} of size {n}"
60    );
61    resolved as usize
62}
63
64/// Multi-axis slice. Apply via [`Tensor::slice`].
65#[derive(Debug, Clone, Default)]
66pub struct SliceSpec {
67    axes: Vec<SliceAxis>,
68}
69
70impl SliceSpec {
71    pub fn new(axes: Vec<SliceAxis>) -> Self {
72        Self { axes }
73    }
74
75    /// Lower to a chain of `narrow` ops (compiler may fuse into consumers).
76    pub fn apply(&self, tensor: &Tensor) -> Tensor {
77        let mut out = tensor.clone();
78        let shape = tensor.shape();
79        for (axis, spec) in self.axes.iter().enumerate() {
80            let dim = match shape.dims().get(axis) {
81                Some(Dim::Static(n)) => Some(*n),
82                Some(Dim::Dynamic(_)) => None,
83                None => panic!("slice axis {axis} out of bounds (rank {})", shape.rank()),
84            };
85            out = match spec {
86                SliceAxis::All => out,
87                SliceAxis::Index(i) => out.narrow(axis, resolve_index(*i, dim, axis), 1),
88                SliceAxis::Range { start, end } => {
89                    let s = resolve_index(*start, dim, axis);
90                    let e = resolve_index(*end, dim, axis);
91                    assert!(e > s, "slice range end must exceed start (axis {axis})");
92                    out.narrow(axis, s, e - s)
93                }
94                SliceAxis::From(start) => {
95                    let s = resolve_index(*start, dim, axis);
96                    let n = dim.unwrap_or_else(|| {
97                        panic!("slice `start..` requires a static dimension on axis {axis}")
98                    });
99                    out.narrow(axis, s, n - s)
100                }
101            };
102        }
103        out
104    }
105}
106
107/// Build a [`SliceSpec`] from axis expressions.
108///
109/// ```rust
110/// use rlx_tensor::{ax, graph, rg, s, shape};
111///
112/// let g = graph("win", |g| {
113///     let x = g.input("x", shape![4, 16]);
114///     x.slice(s![ax(), rg(2, 10)])
115/// });
116/// assert_eq!(g.outputs.len(), 1);
117/// ```
118#[macro_export]
119macro_rules! s {
120    ( $( $axis:expr ),* $(,)? ) => {
121        $crate::SliceSpec::new(vec![ $( $axis ),* ])
122    };
123}
rlx_tensor/slice.rs

rlx_tensor/
slice.rs
