luma_tensor/ops/
arith.rs

1//! `std::ops` trait implementations for [`Tensor`].
2
3use std::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign};
4
5use crate::ops::numeric::NumericDTypeKind;
6use crate::ops::shape::ShapeDTypeKind;
7use crate::{Bool, DTypeKind, Device, Float, Int, Tensor};
8
9// ============================================================================
10//   TensorOrScalar — bridges concrete scalars to K::Scalar (like lumen)
11// ============================================================================
12
13enum TensorOrScalar<D: Device, K: DTypeKind<D>> {
14    Tensor(Tensor<D, K>),
15    Scalar(K::Scalar),
16}
17
18// Concrete scalar → TensorOrScalar (K::Scalar resolves here, not at trait boundary)
19impl<D: Device> From<f64> for TensorOrScalar<D, Float> {
20    fn from(v: f64) -> Self {
21        Self::Scalar(v)
22    }
23}
24impl<D: Device> From<i64> for TensorOrScalar<D, Int> {
25    fn from(v: i64) -> Self {
26        Self::Scalar(v)
27    }
28}
29impl<D: Device> From<bool> for TensorOrScalar<D, Bool> {
30    fn from(v: bool) -> Self {
31        Self::Scalar(v)
32    }
33}
34impl<D: Device, K: DTypeKind<D>> From<Tensor<D, K>> for TensorOrScalar<D, K> {
35    fn from(t: Tensor<D, K>) -> Self {
36        Self::Tensor(t)
37    }
38}
39impl<D: Device, K: DTypeKind<D>> From<&Tensor<D, K>> for TensorOrScalar<D, K> {
40    fn from(t: &Tensor<D, K>) -> Self {
41        Self::Tensor(t.clone())
42    }
43}
44
45// ============================================================================
46//   &t OP rhs   (rhs: impl Into<TensorOrScalar<D, K>>)
47// ============================================================================
48
49macro_rules! impl_ref_op {
50    ($Trait:ident, $method:ident, $scalar_method:ident) => {
51        impl<D, K, R> $Trait<R> for &Tensor<D, K>
52        where
53            D: Device,
54            K: NumericDTypeKind<D> + ShapeDTypeKind<D>,
55            R: Into<TensorOrScalar<D, K>>,
56        {
57            type Output = Tensor<D, K>;
58            fn $method(self, rhs: R) -> Self::Output {
59                match rhs.into() {
60                    TensorOrScalar::Tensor(t) => Tensor::$method(self, &t).unwrap(),
61                    TensorOrScalar::Scalar(s) => Tensor::$scalar_method(self, s).unwrap(),
62                }
63            }
64        }
65
66        impl<D, K, R> $Trait<R> for Tensor<D, K>
67        where
68            D: Device,
69            K: NumericDTypeKind<D> + ShapeDTypeKind<D>,
70            R: Into<TensorOrScalar<D, K>>,
71        {
72            type Output = Tensor<D, K>;
73            fn $method(self, rhs: R) -> Self::Output {
74                match rhs.into() {
75                    TensorOrScalar::Tensor(t) => Tensor::$method(&self, &t).unwrap(),
76                    TensorOrScalar::Scalar(s) => Tensor::$scalar_method(&self, s).unwrap(),
77                }
78            }
79        }
80    };
81}
82
83impl_ref_op!(Add, add, add_scalar);
84impl_ref_op!(Sub, sub, sub_scalar);
85impl_ref_op!(Mul, mul, mul_scalar);
86impl_ref_op!(Div, div, div_scalar);
87
88// ============================================================================
89//   s OP &t   (scalar-left — concrete per scalar type)
90// ============================================================================
91
92impl<D: Device> Add<&Tensor<D, Float>> for f64 {
93    type Output = Tensor<D, Float>;
94    fn add(self, rhs: &Tensor<D, Float>) -> Self::Output {
95        Tensor::add_scalar(rhs, self).unwrap()
96    }
97}
98impl<D: Device> Sub<&Tensor<D, Float>> for f64 {
99    type Output = Tensor<D, Float>;
100    fn sub(self, rhs: &Tensor<D, Float>) -> Self::Output {
101        Tensor::sub_scalar_lhs(rhs, self).unwrap()
102    }
103}
104impl<D: Device> Mul<&Tensor<D, Float>> for f64 {
105    type Output = Tensor<D, Float>;
106    fn mul(self, rhs: &Tensor<D, Float>) -> Self::Output {
107        Tensor::mul_scalar(rhs, self).unwrap()
108    }
109}
110impl<D: Device> Div<&Tensor<D, Float>> for f64 {
111    type Output = Tensor<D, Float>;
112    fn div(self, rhs: &Tensor<D, Float>) -> Self::Output {
113        Tensor::div_scalar_lhs(rhs, self).unwrap()
114    }
115}
116
117impl<D: Device> Add<Tensor<D, Float>> for f64 {
118    type Output = Tensor<D, Float>;
119    fn add(self, rhs: Tensor<D, Float>) -> Self::Output {
120        Tensor::add_scalar(&rhs, self).unwrap()
121    }
122}
123impl<D: Device> Sub<Tensor<D, Float>> for f64 {
124    type Output = Tensor<D, Float>;
125    fn sub(self, rhs: Tensor<D, Float>) -> Self::Output {
126        Tensor::sub_scalar_lhs(&rhs, self).unwrap()
127    }
128}
129impl<D: Device> Mul<Tensor<D, Float>> for f64 {
130    type Output = Tensor<D, Float>;
131    fn mul(self, rhs: Tensor<D, Float>) -> Self::Output {
132        Tensor::mul_scalar(&rhs, self).unwrap()
133    }
134}
135impl<D: Device> Div<Tensor<D, Float>> for f64 {
136    type Output = Tensor<D, Float>;
137    fn div(self, rhs: Tensor<D, Float>) -> Self::Output {
138        Tensor::div_scalar_lhs(&rhs, self).unwrap()
139    }
140}
141
142impl<D: Device> Add<&Tensor<D, Int>> for i64 {
143    type Output = Tensor<D, Int>;
144    fn add(self, rhs: &Tensor<D, Int>) -> Self::Output {
145        Tensor::add_scalar(rhs, self).unwrap()
146    }
147}
148impl<D: Device> Mul<&Tensor<D, Int>> for i64 {
149    type Output = Tensor<D, Int>;
150    fn mul(self, rhs: &Tensor<D, Int>) -> Self::Output {
151        Tensor::mul_scalar(rhs, self).unwrap()
152    }
153}
154
155impl<D: Device> Add<Tensor<D, Int>> for i64 {
156    type Output = Tensor<D, Int>;
157    fn add(self, rhs: Tensor<D, Int>) -> Self::Output {
158        Tensor::add_scalar(&rhs, self).unwrap()
159    }
160}
161impl<D: Device> Mul<Tensor<D, Int>> for i64 {
162    type Output = Tensor<D, Int>;
163    fn mul(self, rhs: Tensor<D, Int>) -> Self::Output {
164        Tensor::mul_scalar(&rhs, self).unwrap()
165    }
166}
167
168// ============================================================================
169//   t OP= &t  +  -&t, -t
170// ============================================================================
171
172macro_rules! impl_assign_and_neg {
173    ($kind:ty) => {
174        impl<D: Device> AddAssign<&Tensor<D, $kind>> for Tensor<D, $kind>
175        where
176            $kind: NumericDTypeKind<D> + ShapeDTypeKind<D>,
177        {
178            fn add_assign(&mut self, rhs: &Tensor<D, $kind>) {
179                Tensor::add_(self, rhs).unwrap()
180            }
181        }
182        impl<D: Device> SubAssign<&Tensor<D, $kind>> for Tensor<D, $kind>
183        where
184            $kind: NumericDTypeKind<D> + ShapeDTypeKind<D>,
185        {
186            fn sub_assign(&mut self, rhs: &Tensor<D, $kind>) {
187                Tensor::sub_(self, rhs).unwrap()
188            }
189        }
190        impl<D: Device> MulAssign<&Tensor<D, $kind>> for Tensor<D, $kind>
191        where
192            $kind: NumericDTypeKind<D> + ShapeDTypeKind<D>,
193        {
194            fn mul_assign(&mut self, rhs: &Tensor<D, $kind>) {
195                Tensor::mul_(self, rhs).unwrap()
196            }
197        }
198        impl<D: Device> DivAssign<&Tensor<D, $kind>> for Tensor<D, $kind>
199        where
200            $kind: NumericDTypeKind<D> + ShapeDTypeKind<D>,
201        {
202            fn div_assign(&mut self, rhs: &Tensor<D, $kind>) {
203                Tensor::div_(self, rhs).unwrap()
204            }
205        }
206
207        impl<D: Device> Neg for &Tensor<D, $kind>
208        where
209            $kind: NumericDTypeKind<D> + ShapeDTypeKind<D>,
210        {
211            type Output = Tensor<D, $kind>;
212            fn neg(self) -> Self::Output {
213                Tensor::neg(self).unwrap()
214            }
215        }
216        impl<D: Device> Neg for Tensor<D, $kind>
217        where
218            $kind: NumericDTypeKind<D> + ShapeDTypeKind<D>,
219        {
220            type Output = Tensor<D, $kind>;
221            fn neg(self) -> Self::Output {
222                -&self
223            }
224        }
225    };
226}
227
228impl_assign_and_neg!(Float);
229impl_assign_and_neg!(Int);
luma_tensor/ops/arith.rs

luma_tensor/ops/
arith.rs
