luma_tensor/ops/
numeric.rs

1use super::shape::ShapeDTypeKind;
2use crate::{BinaryOp, Bool, CmpOp, DTypeKind, Device, Float, FloatMeta, FloatUnaryOp, Int, Layout, Storage, Tensor, TensorMeta, UnaryOp};
3
4// ============================================================================
5//    Numeric dispatch
6// ============================================================================
7
8pub trait NumericDTypeKind<D: Device>: DTypeKind<D> + Sized {
9    fn binary_dispatch(
10        lhs: &Self::Storage,
11        lhs_l: &Layout,
12        rhs: &Self::Storage,
13        rhs_l: &Layout,
14        op: BinaryOp,
15    ) -> crate::Result<Self::Storage>;
16    fn binary_inplace_dispatch(
17        dst: &mut Self::Storage,
18        dst_l: &Layout,
19        src: &Self::Storage,
20        src_l: &Layout,
21        op: BinaryOp,
22    ) -> crate::Result<()>;
23    fn binary_scalar_dispatch(lhs: &Self::Storage, lhs_l: &Layout, rhs: Self::Scalar, op: BinaryOp) -> crate::Result<Self::Storage>;
24    fn binary_scalar_inplace_dispatch(dst: &mut Self::Storage, dst_l: &Layout, rhs: Self::Scalar, op: BinaryOp) -> crate::Result<()>;
25    fn binary_scalar_lhs_dispatch(scalar: Self::Scalar, rhs: &Self::Storage, rhs_l: &Layout, op: BinaryOp) -> crate::Result<Self::Storage>;
26
27    fn cmp_dispatch(lhs: &Self::Storage, lhs_l: &Layout, rhs: &Self::Storage, rhs_l: &Layout, op: CmpOp) -> crate::Result<D::BoolStorage>;
28    fn cmp_scalar_dispatch(lhs: &Self::Storage, lhs_l: &Layout, rhs: Self::Scalar, op: CmpOp) -> crate::Result<D::BoolStorage>;
29
30    fn unary_dispatch(x: &Self::Storage, l: &Layout, op: UnaryOp<Self::Scalar>) -> crate::Result<Self::Storage>;
31    fn unary_inplace_dispatch(dst: &mut Self::Storage, dst_l: &Layout, op: UnaryOp<Self::Scalar>) -> crate::Result<()>;
32}
33
34impl<D: Device> NumericDTypeKind<D> for Float {
35    fn binary_dispatch(
36        lhs: &Self::Storage,
37        lhs_l: &Layout,
38        rhs: &Self::Storage,
39        rhs_l: &Layout,
40        op: BinaryOp,
41    ) -> crate::Result<Self::Storage> {
42        D::f_binary(lhs, lhs_l, rhs, rhs_l, op)
43    }
44    fn binary_inplace_dispatch(
45        dst: &mut Self::Storage,
46        dst_l: &Layout,
47        src: &Self::Storage,
48        src_l: &Layout,
49        op: BinaryOp,
50    ) -> crate::Result<()> {
51        D::f_binary_(dst, dst_l, src, src_l, op)
52    }
53    fn binary_scalar_dispatch(lhs: &Self::Storage, lhs_l: &Layout, rhs: Self::Scalar, op: BinaryOp) -> crate::Result<Self::Storage> {
54        D::f_binary_scalar(lhs, lhs_l, rhs, op)
55    }
56    fn binary_scalar_inplace_dispatch(dst: &mut Self::Storage, dst_l: &Layout, rhs: Self::Scalar, op: BinaryOp) -> crate::Result<()> {
57        D::f_binary_scalar_(dst, dst_l, rhs, op)
58    }
59    fn binary_scalar_lhs_dispatch(scalar: Self::Scalar, rhs: &Self::Storage, rhs_l: &Layout, op: BinaryOp) -> crate::Result<Self::Storage> {
60        D::f_binary_scalar_lhs(scalar, rhs, rhs_l, op)
61    }
62    fn cmp_dispatch(lhs: &Self::Storage, lhs_l: &Layout, rhs: &Self::Storage, rhs_l: &Layout, op: CmpOp) -> crate::Result<D::BoolStorage> {
63        D::f_cmp(lhs, lhs_l, rhs, rhs_l, op)
64    }
65    fn cmp_scalar_dispatch(lhs: &Self::Storage, lhs_l: &Layout, rhs: Self::Scalar, op: CmpOp) -> crate::Result<D::BoolStorage> {
66        D::f_cmp_scalar(lhs, lhs_l, rhs, op)
67    }
68    fn unary_dispatch(x: &Self::Storage, l: &Layout, op: UnaryOp<Self::Scalar>) -> crate::Result<Self::Storage> {
69        D::f_unary(x, l, op)
70    }
71    fn unary_inplace_dispatch(dst: &mut Self::Storage, dst_l: &Layout, op: UnaryOp<Self::Scalar>) -> crate::Result<()> {
72        D::f_unary_(dst, dst_l, op)
73    }
74}
75
76impl<D: Device> NumericDTypeKind<D> for Int {
77    fn binary_dispatch(
78        lhs: &Self::Storage,
79        lhs_l: &Layout,
80        rhs: &Self::Storage,
81        rhs_l: &Layout,
82        op: BinaryOp,
83    ) -> crate::Result<Self::Storage> {
84        D::i_binary(lhs, lhs_l, rhs, rhs_l, op)
85    }
86    fn binary_inplace_dispatch(
87        dst: &mut Self::Storage,
88        dst_l: &Layout,
89        src: &Self::Storage,
90        src_l: &Layout,
91        op: BinaryOp,
92    ) -> crate::Result<()> {
93        D::i_binary_(dst, dst_l, src, src_l, op)
94    }
95    fn binary_scalar_dispatch(lhs: &Self::Storage, lhs_l: &Layout, rhs: Self::Scalar, op: BinaryOp) -> crate::Result<Self::Storage> {
96        D::i_binary_scalar(lhs, lhs_l, rhs, op)
97    }
98    fn binary_scalar_inplace_dispatch(dst: &mut Self::Storage, dst_l: &Layout, rhs: Self::Scalar, op: BinaryOp) -> crate::Result<()> {
99        D::i_binary_scalar_(dst, dst_l, rhs, op)
100    }
101    fn binary_scalar_lhs_dispatch(scalar: Self::Scalar, rhs: &Self::Storage, rhs_l: &Layout, op: BinaryOp) -> crate::Result<Self::Storage> {
102        D::i_binary_scalar_lhs(scalar, rhs, rhs_l, op)
103    }
104    fn cmp_dispatch(lhs: &Self::Storage, lhs_l: &Layout, rhs: &Self::Storage, rhs_l: &Layout, op: CmpOp) -> crate::Result<D::BoolStorage> {
105        D::i_cmp(lhs, lhs_l, rhs, rhs_l, op)
106    }
107    fn cmp_scalar_dispatch(lhs: &Self::Storage, lhs_l: &Layout, rhs: Self::Scalar, op: CmpOp) -> crate::Result<D::BoolStorage> {
108        D::i_cmp_scalar(lhs, lhs_l, rhs, op)
109    }
110    fn unary_dispatch(x: &Self::Storage, l: &Layout, op: UnaryOp<Self::Scalar>) -> crate::Result<Self::Storage> {
111        D::i_unary(x, l, op)
112    }
113    fn unary_inplace_dispatch(dst: &mut Self::Storage, dst_l: &Layout, op: UnaryOp<Self::Scalar>) -> crate::Result<()> {
114        D::i_unary_(dst, dst_l, op)
115    }
116}
117
118// ============================================================================
119//   impl binary op
120// ============================================================================
121
122impl<D: Device, K: NumericDTypeKind<D> + ShapeDTypeKind<D>> Tensor<D, K> {
123    pub(crate) fn binary_impl(&self, rhs: &Self, op: BinaryOp, name: &'static str) -> crate::Result<Self> {
124        let shape = self.same_shape(rhs, name)?.clone();
125        let s = K::binary_dispatch(&*self.storage_read()?, self.layout(), &*rhs.storage_read()?, rhs.layout(), op)?;
126        assert_eq!(self.dtype(), s.dtype());
127        Ok(Self::from_storage(s, shape, K::Meta::on_binary(self, rhs, op)))
128    }
129
130    pub(crate) fn binary_scalar_impl(&self, rhs: K::Scalar, op: BinaryOp) -> crate::Result<Self> {
131        let storage = K::binary_scalar_dispatch(&*self.storage_read()?, self.layout(), rhs, op)?;
132        let meta = K::Meta::on_binary_scalar(self, rhs, op);
133        assert_eq!(self.dtype(), storage.dtype());
134        Ok(Self::from_storage(storage, self.shape().clone(), meta))
135    }
136
137    pub(crate) fn binary_inplace_impl(&self, rhs: &Self, op: BinaryOp) -> crate::Result<()> {
138        let mut s = self.storage_write()?;
139        K::binary_inplace_dispatch(&mut s, self.layout(), &*rhs.storage_read()?, rhs.layout(), op)
140    }
141
142    pub(crate) fn binary_scalar_inplace_impl(&self, rhs: K::Scalar, op: BinaryOp) -> crate::Result<()> {
143        let mut s = self.storage_write()?;
144        K::binary_scalar_inplace_dispatch(&mut s, self.layout(), rhs, op)
145    }
146
147    pub(crate) fn binary_broadcast_impl(&self, rhs: &Self, op: BinaryOp, name: &'static str) -> crate::Result<Self> {
148        let out_shape = self.shape().broadcast_shape_binary_op(rhs.shape(), name)?;
149        let lhs = self.broadcast_as(out_shape.clone())?;
150        let rhs = rhs.broadcast_as(out_shape)?;
151        lhs.binary_impl(&rhs, op, name)
152    }
153}
154
155macro_rules! binary_impl {
156    ($name:ident, $variant:ident) => {
157        paste::paste! {
158            #[inline]
159            pub fn $name(&self, rhs: &Self) -> crate::Result<Self> {
160                self.binary_impl(rhs, BinaryOp::$variant, stringify!($name))
161            }
162
163            #[inline]
164            pub fn [<$name _scalar>](&self, rhs: K::Scalar) -> crate::Result<Self> {
165                self.binary_scalar_impl(rhs, BinaryOp::$variant)
166            }
167
168            #[inline]
169            pub fn [<$name _>](&self, rhs: &Self) -> crate::Result<()> {
170                self.binary_inplace_impl(rhs, BinaryOp::$variant)
171            }
172
173            #[inline]
174            pub fn [<$name _scalar_>](&self, rhs: K::Scalar) -> crate::Result<()> {
175                self.binary_scalar_inplace_impl(rhs, BinaryOp::$variant)
176            }
177
178            #[inline]
179            pub fn [<broadcast_ $name>](&self, rhs: &Self) -> crate::Result<Self> {
180                self.binary_broadcast_impl(rhs, BinaryOp::$variant, stringify!([<broadcast_ $name>]))
181            }
182        }
183    };
184}
185
186impl<D: Device, K: NumericDTypeKind<D> + ShapeDTypeKind<D>> Tensor<D, K> {
187    binary_impl!(add, Add);
188    binary_impl!(sub, Sub);
189    binary_impl!(mul, Mul);
190    binary_impl!(div, Div);
191    binary_impl!(maximum, Maximum);
192    binary_impl!(minimum, Minimum);
193}
194
195// ============================================================================
196//   impl cmp op
197// ============================================================================
198
199impl<D: Device, K: NumericDTypeKind<D> + ShapeDTypeKind<D>> Tensor<D, K> {
200    pub(crate) fn cmp_impl(&self, rhs: &Self, op: CmpOp, name: &'static str) -> crate::Result<Tensor<D, Bool>> {
201        let shape = self.same_shape(rhs, name)?.clone();
202        let s = K::cmp_dispatch(&*self.storage_read()?, self.layout(), &*rhs.storage_read()?, rhs.layout(), op)?;
203        Ok(Tensor::<D, Bool>::from_storage(s, shape, ()))
204    }
205
206    pub(crate) fn cmp_scalar_imlp(&self, rhs: K::Scalar, op: CmpOp) -> crate::Result<Tensor<D, Bool>> {
207        let s = K::cmp_scalar_dispatch(&*self.storage_read()?, self.layout(), rhs, op)?;
208        Ok(Tensor::<D, Bool>::from_storage(s, self.shape(), ()))
209    }
210
211    pub(crate) fn broadcast_cmp_impl(&self, rhs: &Self, op: CmpOp, name: &'static str) -> crate::Result<Tensor<D, Bool>> {
212        let out_shape = self.shape().broadcast_shape_binary_op(rhs.shape(), name)?;
213        let lhs = self.broadcast_as(out_shape.clone())?;
214        let rhs = rhs.broadcast_as(out_shape)?;
215        lhs.cmp_impl(&rhs, op, name)
216    }
217}
218
219macro_rules! cmp_impl {
220    ($name:ident, $variant:ident) => {
221        paste::paste! {
222            #[inline]
223            pub fn $name(&self, rhs: &Self) -> crate::Result<Tensor<D, Bool>> {
224                self.cmp_impl(rhs, CmpOp::$variant, stringify!($name))
225            }
226
227            #[inline]
228            pub fn [<$name _scalar>](&self, rhs: K::Scalar) -> crate::Result<Tensor<D, Bool>> {
229                self.cmp_scalar_imlp(rhs, CmpOp::$variant)
230            }
231
232            #[inline]
233            pub fn [<broadcast_ $name>](&self, rhs: &Self) -> crate::Result<Tensor<D, Bool>> {
234                self.broadcast_cmp_impl(rhs, CmpOp::$variant, stringify!([<broadcast_ $name>]))
235            }
236        }
237    };
238}
239
240impl<D: Device, K: NumericDTypeKind<D> + ShapeDTypeKind<D>> Tensor<D, K> {
241    cmp_impl!(eq, Eq);
242    cmp_impl!(ne, Ne);
243    cmp_impl!(lt, Lt);
244    cmp_impl!(le, Le);
245    cmp_impl!(gt, Gt);
246    cmp_impl!(ge, Ge);
247}
248
249// ============================================================================
250//   unary
251// ============================================================================
252
253impl<D: Device, K: NumericDTypeKind<D> + ShapeDTypeKind<D>> Tensor<D, K> {
254    pub(crate) fn unary_impl(&self, op: UnaryOp<K::Scalar>) -> crate::Result<Tensor<D, K>> {
255        let storage = K::unary_dispatch(&*self.storage_read()?, self.layout(), op)?;
256        let meta = K::Meta::on_unary(self, op);
257        Ok(Self::from_storage(storage, self.shape().clone(), meta))
258    }
259
260    pub(crate) fn unary_inplace_impl(&self, op: UnaryOp<K::Scalar>) -> crate::Result<()> {
261        let mut s = self.storage_write()?;
262        K::unary_inplace_dispatch(&mut s, self.layout(), op)
263    }
264}
265
266macro_rules! unary_method {
267    ($name:ident, $op:tt) => {
268        paste::paste! {
269            #[inline]
270            pub fn $name(&self) -> crate::Result<Self> {
271                self.unary_impl(UnaryOp::$op)
272            }
273
274            #[inline]
275            pub fn [<$name _>](&self) -> crate::Result<()> {
276                self.unary_inplace_impl(UnaryOp::$op)
277            }
278        }
279    };
280}
281
282impl<D: Device, K: NumericDTypeKind<D> + ShapeDTypeKind<D>> Tensor<D, K> {
283    unary_method!(neg, Neg);
284    unary_method!(abs, Abs);
285    unary_method!(sign, Sign);
286
287    #[inline]
288    pub fn affine(&self, mul: K::Scalar, add: K::Scalar) -> crate::Result<Self> {
289        self.unary_impl(UnaryOp::Affine(mul, add))
290    }
291
292    #[inline]
293    pub fn affine_(&self, mul: K::Scalar, add: K::Scalar) -> crate::Result<()> {
294        self.unary_inplace_impl(UnaryOp::Affine(mul, add))
295    }
296
297    #[inline]
298    pub fn pow(&self, exp: K::Scalar) -> crate::Result<Self> {
299        self.unary_impl(UnaryOp::Pow(exp))
300    }
301
302    #[inline]
303    pub fn pow_(&self, exp: K::Scalar) -> crate::Result<()> {
304        self.unary_inplace_impl(UnaryOp::Pow(exp))
305    }
306
307    #[inline]
308    pub fn clamp(&self, min: Option<K::Scalar>, max: Option<K::Scalar>) -> crate::Result<Self> {
309        self.unary_impl(UnaryOp::Clamp(min, max))
310    }
311
312    #[inline]
313    pub fn clamp_(&self, min: Option<K::Scalar>, max: Option<K::Scalar>) -> crate::Result<()> {
314        self.unary_inplace_impl(UnaryOp::Clamp(min, max))
315    }
316
317    /// `scalar - self` (scalar on the **left**).
318    #[inline]
319    pub fn sub_scalar_lhs(&self, lhs: K::Scalar) -> crate::Result<Self> {
320        let storage = K::binary_scalar_lhs_dispatch(lhs, &*self.storage_read()?, self.layout(), BinaryOp::Sub)?;
321        let meta = K::Meta::default();
322        Ok(Self::from_storage(storage, self.shape().clone(), meta))
323    }
324
325    /// `scalar / self` (scalar on the **left**).
326    #[inline]
327    pub fn div_scalar_lhs(&self, lhs: K::Scalar) -> crate::Result<Self> {
328        let storage = K::binary_scalar_lhs_dispatch(lhs, &*self.storage_read()?, self.layout(), BinaryOp::Div)?;
329        let meta = K::Meta::default();
330        Ok(Self::from_storage(storage, self.shape().clone(), meta))
331    }
332}
333
334// ============================================================================
335//   activate for float
336// ============================================================================
337
338macro_rules! float_unary_method {
339    ($name:ident, $variant:ident) => {
340        paste::paste! {
341            #[inline]
342            pub fn $name(&self) -> crate::Result<Self> {
343                self.float_unary_impl(FloatUnaryOp::$variant)
344            }
345
346            #[inline]
347            pub fn [<$name _>](&self) -> crate::Result<()> {
348                self.float_unary_inplace_impl(FloatUnaryOp::$variant)
349            }
350        }
351    };
352}
353
354impl<D: Device> Tensor<D, Float> {
355    fn float_unary_impl(&self, op: FloatUnaryOp) -> crate::Result<Self> {
356        let storage = D::f_float_unary(&*self.storage_read()?, self.layout(), op)?;
357        let meta = FloatMeta::on_float_unary(self, op);
358        Ok(Self::from_storage(storage, self.shape().clone(), meta))
359    }
360
361    fn float_unary_inplace_impl(&self, op: FloatUnaryOp) -> crate::Result<()> {
362        let mut s = self.storage_write()?;
363        D::f_float_unary_(&mut s, self.layout(), op)
364    }
365
366    float_unary_method!(exp, Exp);
367    float_unary_method!(ln, Ln);
368    float_unary_method!(sin, Sin);
369    float_unary_method!(cos, Cos);
370    float_unary_method!(tanh, Tanh);
371    float_unary_method!(sqr, Sqr);
372    float_unary_method!(sqrt, Sqrt);
373    float_unary_method!(recip, Recip);
374    float_unary_method!(gelu, Gelu);
375    float_unary_method!(gelu_erf, GeluErf);
376    float_unary_method!(erf, Erf);
377    float_unary_method!(relu, Relu);
378    float_unary_method!(silu, Silu);
379    float_unary_method!(sigmoid, Sigmoid);
380    float_unary_method!(floor, Floor);
381    float_unary_method!(ceil, Ceil);
382    float_unary_method!(round, Round);
383
384    pub fn leaky_relu(&self, negative_slope: f64) -> crate::Result<Self> {
385        self.float_unary_impl(FloatUnaryOp::LeakyRelu(negative_slope))
386    }
387
388    pub fn leaky_relu_(&self, negative_slope: f64) -> crate::Result<()> {
389        self.float_unary_inplace_impl(FloatUnaryOp::LeakyRelu(negative_slope))
390    }
391}
luma_tensor/ops/numeric.rs

luma_tensor/ops/
numeric.rs
