luma_tensor/device/
int_ops.rs

1use std::borrow::Cow;
2
3use crate::Result;
4use crate::dtype::{BoolDType, FloatDType, IntDType};
5use crate::tensor::{Layout, Shape};
6
7/// Operations on `Int`-kind tensors. No autograd. Scalars are `i64` (wide enough
8/// for all int precisions).
9pub trait IntOps<D: super::Device> {
10    // construction
11    fn i_zeros(shape: &Shape, device: &D, dtype: IntDType) -> Result<D::IntStorage>;
12    fn i_ones(shape: &Shape, device: &D, dtype: IntDType) -> Result<D::IntStorage>;
13    fn i_full(shape: &Shape, value: i64, device: &D, dtype: IntDType) -> Result<D::IntStorage>;
14    fn i_from_i64<'a>(data: impl Into<Cow<'a, [i64]>>, device: &D) -> Result<D::IntStorage>;
15
16    fn i_from_i32<'a>(data: impl Into<Cow<'a, [i32]>>, device: &D) -> Result<D::IntStorage>;
17
18    fn i_from_u32<'a>(data: impl Into<Cow<'a, [u32]>>, device: &D) -> Result<D::IntStorage>;
19
20    fn i_from_u8<'a>(data: impl Into<Cow<'a, [u8]>>, device: &D) -> Result<D::IntStorage>;
21
22    fn i_from_bytes<'a>(bytes: impl Into<Cow<'a, [u8]>>, shape: &Shape, device: &D, dtype: IntDType) -> Result<D::IntStorage>;
23
24    fn i_arange(start: i64, end: i64, step: i64, device: &D, dtype: IntDType) -> Result<(D::IntStorage, usize)>;
25
26    // materialization / read-back
27    fn i_contiguous(x: &D::IntStorage, layout: &Layout) -> Result<D::IntStorage>;
28    fn i_cast_float(x: &D::IntStorage, layout: &Layout, to: FloatDType) -> Result<D::FloatStorage>;
29    fn i_cast_int(x: &D::IntStorage, layout: &Layout, to: IntDType) -> Result<D::IntStorage>;
30    fn i_cast_bool(x: &D::IntStorage, layout: &Layout, to: BoolDType) -> Result<D::BoolStorage>;
31
32    /// Read all elements into a `Vec<i64>` in logical (layout) order.
33    fn i_to_vec(x: &D::IntStorage, layout: &Layout) -> Result<Vec<i64>>;
34
35    /// Read raw little-endian bytes in logical (layout) order.
36    /// Returns `Cow::Borrowed` when the underlying storage is already contiguous
37    /// (zero-copy); `Cow::Owned` otherwise.
38    fn i_to_bytes<'a>(x: &'a D::IntStorage, layout: &Layout) -> Result<Cow<'a, [u8]>>;
39
40    // arithmetic
41    fn i_binary(lhs: &D::IntStorage, lhs_l: &Layout, rhs: &D::IntStorage, rhs_l: &Layout, op: crate::BinaryOp) -> Result<D::IntStorage>;
42
43    fn i_binary_(dst: &mut D::IntStorage, dst_l: &Layout, src: &D::IntStorage, src_l: &Layout, op: crate::BinaryOp) -> Result<()>;
44
45    fn i_binary_scalar(lhs: &D::IntStorage, lhs_l: &Layout, rhs: i64, op: crate::BinaryOp) -> Result<D::IntStorage>;
46
47    fn i_binary_scalar_(dst: &mut D::IntStorage, dst_l: &Layout, rhs: i64, op: crate::BinaryOp) -> Result<()>;
48
49    fn i_binary_scalar_lhs(scalar: i64, rhs: &D::IntStorage, rhs_l: &Layout, op: crate::BinaryOp) -> Result<D::IntStorage>;
50
51    fn i_unary(x: &D::IntStorage, layout: &Layout, op: crate::UnaryOp<i64>) -> Result<D::IntStorage>;
52
53    fn i_unary_(dst: &mut D::IntStorage, dst_l: &Layout, op: crate::UnaryOp<i64>) -> Result<()>;
54
55    // matmul (batched); out shape computed by the caller / this fn
56    fn i_matmul(lhs: &D::IntStorage, lhs_l: &Layout, rhs: &D::IntStorage, rhs_l: &Layout) -> Result<(D::IntStorage, Shape)>;
57
58    // comparison -> bool
59    fn i_cmp(lhs: &D::IntStorage, lhs_l: &Layout, rhs: &D::IntStorage, rhs_l: &Layout, op: crate::CmpOp) -> Result<D::BoolStorage>;
60
61    fn i_cmp_scalar(lhs: &D::IntStorage, lhs_l: &Layout, rhs: i64, op: crate::CmpOp) -> Result<D::BoolStorage>;
62
63    // reduction
64    fn i_reduce(x: &D::IntStorage, layout: &Layout, dims: &[usize], keepdim: bool, op: crate::ReduceOp) -> Result<(D::IntStorage, Shape)>;
65
66    /// argmin/argmax return int-kind indices.
67    fn i_arg_reduce(x: &D::IntStorage, layout: &Layout, dim: usize, keepdim: bool, take_max: bool) -> Result<(D::IntStorage, Shape)>;
68
69    // indexing
70    fn i_index_select(x: &D::IntStorage, x_l: &Layout, idx: &D::IntStorage, idx_l: &Layout, dim: usize) -> Result<(D::IntStorage, Shape)>;
71
72    fn i_gather(x: &D::IntStorage, x_l: &Layout, idx: &D::IntStorage, idx_l: &Layout, dim: usize) -> Result<(D::IntStorage, Shape)>;
73
74    fn i_index_add(
75        init: &D::IntStorage,
76        init_l: &Layout,
77        idx: &D::IntStorage,
78        idx_l: &Layout,
79        src: &D::IntStorage,
80        src_l: &Layout,
81        dim: usize,
82    ) -> Result<D::IntStorage>;
83
84    fn i_scatter_add(
85        init: &D::IntStorage,
86        init_l: &Layout,
87        idx: &D::IntStorage,
88        idx_l: &Layout,
89        src: &D::IntStorage,
90        src_l: &Layout,
91        dim: usize,
92    ) -> Result<D::IntStorage>;
93
94    // shape
95    fn i_cat(srcs: &[(&D::IntStorage, &Layout)], dim: usize) -> Result<(D::IntStorage, Shape)>;
96
97    /// Produce the storage for a view of `src` under `dst_l`. See [`FloatOps::f_view`].
98    fn i_view(_src: &D::IntStorage, _src_l: &Layout, _dst_l: &Layout, _view: crate::ViewOp) -> Result<Option<D::IntStorage>> {
99        Ok(None)
100    }
101
102    // pick via a bool mask
103    fn i_pick(
104        mask: &D::BoolStorage,
105        mask_l: &Layout,
106        on_true: &D::IntStorage,
107        true_l: &Layout,
108        on_false: &D::IntStorage,
109        false_l: &Layout,
110    ) -> Result<D::IntStorage>;
111
112    fn i_pick_true(mask: &D::BoolStorage, mask_l: &Layout, value: i64, on_false: &D::IntStorage, false_l: &Layout)
113    -> Result<D::IntStorage>;
114
115    fn i_pick_false(mask: &D::BoolStorage, mask_l: &Layout, on_true: &D::IntStorage, true_l: &Layout, value: i64) -> Result<D::IntStorage>;
116
117    // allclose
118    fn i_allclose(a: &D::IntStorage, a_l: &Layout, b: &D::IntStorage, b_l: &Layout) -> Result<bool>;
119}
luma_tensor/device/int_ops.rs

luma_tensor/device/
int_ops.rs
