Skip to main content

r2l_core/tensor/
candle_tensor.rs

1use candle_core::{Device, Tensor};
2use itertools::izip;
3
4use crate::{
5    error::{Error, TensorError},
6    tensor::{R2lTensor, VecTensor},
7};
8
9type Result<T> = std::result::Result<T, TensorError>;
10
11impl From<candle_core::Error> for Error {
12    fn from(error: candle_core::Error) -> Self {
13        TensorError::operation("Candle backend", error).into()
14    }
15}
16
17impl R2lTensor for Tensor {
18    fn to_vec(&self) -> Result<Vec<f32>> {
19        self.flatten_all()
20            .and_then(|tensor| tensor.to_vec1())
21            .map_err(|error| TensorError::operation("convert to vector", error))
22    }
23
24    fn to_shape(&self) -> Vec<usize> {
25        self.shape().dims().to_vec()
26    }
27
28    fn from_slice_and_shape(data: &[f32], shape: Vec<usize>) -> Result<Self> {
29        validate_shape(data.len(), &shape)?;
30        Tensor::from_slice(data, shape, &Device::Cpu)
31            .map_err(|error| TensorError::operation("construct from slice", error))
32    }
33
34    fn from_vec_and_shape(data: Vec<f32>, shape: Vec<usize>) -> Result<Self> {
35        validate_shape(data.len(), &shape)?;
36        Tensor::from_vec(data, shape, &Device::Cpu)
37            .map_err(|error| TensorError::operation("construct from vector", error))
38    }
39
40    fn add(&self, other: &Self) -> Result<Self> {
41        ensure_same_shape(self, other, "add")?;
42        self.add(other)
43            .map_err(|error| TensorError::operation("add", error))
44    }
45
46    fn sub(&self, other: &Self) -> Result<Self> {
47        ensure_same_shape(self, other, "subtract")?;
48        self.sub(other)
49            .map_err(|error| TensorError::operation("subtract", error))
50    }
51
52    fn mul(&self, other: &Self) -> Result<Self> {
53        ensure_same_shape(self, other, "multiply")?;
54        self.mul(other)
55            .map_err(|error| TensorError::operation("multiply", error))
56    }
57
58    fn exp(&self) -> Result<Self> {
59        self.exp()
60            .map_err(|error| TensorError::operation("exponential", error))
61    }
62
63    fn clamp(&self, min: f32, max: f32) -> Result<Self> {
64        self.clamp(min, max)
65            .map_err(|error| TensorError::operation("clamp", error))
66    }
67
68    fn minimum(&self, other: &Self) -> Result<Self> {
69        ensure_same_shape(self, other, "minimum")?;
70        Self::minimum(self, other).map_err(|error| TensorError::operation("minimum", error))
71    }
72
73    fn neg(&self) -> Result<Self> {
74        self.neg()
75            .map_err(|error| TensorError::operation("negate", error))
76    }
77
78    fn mean(&self) -> Result<Self> {
79        if self.elem_count() == 0 {
80            return Err(TensorError::EmptyInput {
81                operation: "mean".into(),
82            });
83        }
84        self.mean_all()
85            .map_err(|error| TensorError::operation("mean", error))
86    }
87
88    fn sqr(&self) -> Result<Self> {
89        self.sqr()
90            .map_err(|error| TensorError::operation("square", error))
91    }
92
93    fn zeros(shape: Vec<usize>) -> Result<Self> {
94        Tensor::zeros(shape, candle_core::DType::F32, &Device::Cpu)
95            .map_err(|error| TensorError::operation("create zeros", error))
96    }
97
98    fn mul_scalar(&self, scalar: f32) -> Result<Self> {
99        let scalar = Tensor::full(scalar, (), self.device())
100            .map_err(|error| TensorError::operation("create scalar", error))?;
101        self.broadcast_mul(&scalar)
102            .map_err(|error| TensorError::operation("multiply by scalar", error))
103    }
104}
105
106fn validate_shape(data_len: usize, shape: &[usize]) -> Result<()> {
107    let expected = shape.iter().product();
108    if expected != data_len {
109        return Err(TensorError::InvalidShape {
110            shape: shape.to_vec(),
111            expected,
112            actual: data_len,
113        });
114    }
115    Ok(())
116}
117
118fn ensure_same_shape(left: &Tensor, right: &Tensor, operation: &str) -> Result<()> {
119    let left = left.shape().dims().to_vec();
120    let right = right.shape().dims().to_vec();
121    if left != right {
122        return Err(TensorError::ShapeMismatch {
123            operation: operation.into(),
124            left,
125            right,
126        });
127    }
128    Ok(())
129}
130
131impl VecTensor {
132    /// Clamps each element between the corresponding values in `min` and `max`.
133    ///
134    /// # Errors
135    ///
136    /// Returns an error if the tensor shapes differ or the result cannot be constructed.
137    pub fn clamp(&self, min: &Self, max: &Self) -> Result<Self> {
138        self.ensure_same_shape(min, "clamp minimum")?;
139        self.ensure_same_shape(max, "clamp maximum")?;
140        let data = izip!(&self.data, &min.data, &max.data)
141            .map(|(value, min, max)| value.clamp(*min, *max))
142            .collect();
143        Self::new(data, self.shape.clone())
144    }
145}