all_close/
lib.rs

1use candle_core::{
2    DType, Error, Result, Shape, Tensor,
3    scalar::{TensorOrScalar, TensorScalar},
4};
5
6// https://github.com/huggingface/candle/pull/1549/files
7pub trait TensorAllClose {
8    fn all_close<T: TensorOrScalar>(&self, rhs: T, tolerance: f64) -> Result<bool>;
9}
10
11impl TensorAllClose for Tensor {
12    fn all_close<T: TensorOrScalar>(&self, rhs: T, tolerance: f64) -> Result<bool> {
13        let rhs = match rhs.to_tensor_scalar()? {
14            TensorScalar::Tensor(rhs) => rhs,
15            TensorScalar::Scalar(rhs) => rhs
16                .to_dtype(self.dtype())?
17                .to_device(self.device())?
18                .broadcast_as(self.shape())?,
19        };
20        let shape = same_shape_binary_op(&self, &rhs, "all_close")?;
21        let all = self
22            .sub(&rhs)?
23            .abs()?
24            .le(tolerance)?
25            .to_dtype(DType::U32)?
26            .sum_all()?;
27        Ok(all.to_scalar::<u32>()? == shape.elem_count() as u32)
28    }
29}
30
31pub(crate) fn same_shape_binary_op<'a>(
32    lhs: &'a Tensor,
33    rhs: &Tensor,
34    op: &'static str,
35) -> Result<&'a Shape> {
36    let lhs = lhs.shape();
37    let rhs = rhs.shape();
38    if lhs != rhs {
39        Err(Error::ShapeMismatchBinaryOp {
40            lhs: lhs.clone(),
41            rhs: rhs.clone(),
42            op,
43        }
44        .bt())
45    } else {
46        Ok(lhs)
47    }
48}
49
50#[cfg(test)]
51mod tests {
52    use super::*;
53    use candle_core::{Device, test_device};
54
55    fn all_close(device: &Device) -> Result<()> {
56        let t1 = Tensor::new(&[1.0004_f32, 1.0005], device)?;
57        let t2 = Tensor::new(&[1.0005_f32, 1.0004], device)?;
58        let x = t1.all_close(&t2, 0.001)?;
59        let y = t1.all_close(&t2, 0.00001)?;
60        assert_eq!(x, true);
61        assert_eq!(y, false);
62        Ok(())
63    }
64
65    #[test]
66    fn all_close_works() {
67        let device = Device::Cpu;
68
69        all_close(&device).unwrap();
70    }
71}
all_close/lib.rs

all_close/
lib.rs
