rlx_tensor/
slice.rs

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

rlx_tensor/
slice.rs
