oxmera_tensor/
overload.rs

1//! `std::ops` sugar for tensors.
2//!
3//! Operators must return a value, not a `Result`, so these impls panic on
4//! shape/device/dtype errors. Use the named methods (`Tensor::add`, …)
5//! when the failure should be handled instead of being a bug.
6
7use crate::tensor::Tensor;
8
9macro_rules! binary_operator {
10    ($trait:ident, $method:ident, $tensor_method:ident) => {
11        impl std::ops::$trait<&Tensor> for &Tensor {
12            type Output = Tensor;
13            fn $method(self, rhs: &Tensor) -> Tensor {
14                Tensor::$tensor_method(self, rhs)
15                    .unwrap_or_else(|e| panic!(concat!("tensor ", stringify!($method), ": {}"), e))
16            }
17        }
18        impl std::ops::$trait<Tensor> for Tensor {
19            type Output = Tensor;
20            fn $method(self, rhs: Tensor) -> Tensor {
21                std::ops::$trait::$method(&self, &rhs)
22            }
23        }
24        impl std::ops::$trait<&Tensor> for Tensor {
25            type Output = Tensor;
26            fn $method(self, rhs: &Tensor) -> Tensor {
27                std::ops::$trait::$method(&self, rhs)
28            }
29        }
30        impl std::ops::$trait<Tensor> for &Tensor {
31            type Output = Tensor;
32            fn $method(self, rhs: Tensor) -> Tensor {
33                std::ops::$trait::$method(self, &rhs)
34            }
35        }
36        impl std::ops::$trait<f32> for &Tensor {
37            type Output = Tensor;
38            fn $method(self, rhs: f32) -> Tensor {
39                let rhs =
40                    Tensor::scalar_on(self, rhs).unwrap_or_else(|e| panic!("scalar operand: {e}"));
41                std::ops::$trait::$method(self, &rhs)
42            }
43        }
44        impl std::ops::$trait<f32> for Tensor {
45            type Output = Tensor;
46            fn $method(self, rhs: f32) -> Tensor {
47                std::ops::$trait::$method(&self, rhs)
48            }
49        }
50        impl std::ops::$trait<&Tensor> for f32 {
51            type Output = Tensor;
52            fn $method(self, rhs: &Tensor) -> Tensor {
53                let lhs =
54                    Tensor::scalar_on(rhs, self).unwrap_or_else(|e| panic!("scalar operand: {e}"));
55                std::ops::$trait::$method(&lhs, rhs)
56            }
57        }
58        impl std::ops::$trait<Tensor> for f32 {
59            type Output = Tensor;
60            fn $method(self, rhs: Tensor) -> Tensor {
61                std::ops::$trait::$method(self, &rhs)
62            }
63        }
64    };
65}
66
67binary_operator!(Add, add, add);
68binary_operator!(Sub, sub, sub);
69binary_operator!(Mul, mul, mul);
70binary_operator!(Div, div, div);
71
72impl std::ops::Neg for &Tensor {
73    type Output = Tensor;
74    fn neg(self) -> Tensor {
75        Tensor::neg(self).unwrap_or_else(|e| panic!("tensor neg: {e}"))
76    }
77}
78
79impl std::ops::Neg for Tensor {
80    type Output = Tensor;
81    fn neg(self) -> Tensor {
82        std::ops::Neg::neg(&self)
83    }
84}
oxmera_tensor/overload.rs

oxmera_tensor/
overload.rs
