luma_tensor/ops/
display.rs

1//! Display, to_scalar, allclose, false_count, and other utility ops.
2
3use crate::dtype::{FloatDType, IntDType};
4use crate::{Bool, Device, Float, Int, Tensor};
5
6// ---- to_scalar ----
7
8impl<D: Device> Tensor<D, Float> {
9    /// Read the single element of a scalar (0-d or 1-element) tensor.
10    pub fn to_scalar(&self) -> crate::Result<f64> {
11        if self.element_count() != 1 {
12            return Err(crate::Error::NotScalar);
13        }
14        Ok(self.to_vec()?[0])
15    }
16}
17
18impl<D: Device> Tensor<D, Int> {
19    pub fn to_scalar(&self) -> crate::Result<i64> {
20        if self.element_count() != 1 {
21            return Err(crate::Error::NotScalar);
22        }
23        Ok(self.to_vec()?[0])
24    }
25}
26
27// ---- allclose ----
28
29impl<D: Device> Tensor<D, Float> {
30    pub fn allclose(&self, other: &Self, rtol: f64, atol: f64) -> crate::Result<bool> {
31        if self.element_count() != other.element_count() {
32            return Ok(false);
33        }
34        D::f_allclose(&*self.storage_read()?, self.layout(), &*other.storage_read()?, other.layout(), rtol, atol)
35    }
36}
37
38impl<D: Device> Tensor<D, Int> {
39    pub fn allclose(&self, other: &Self) -> crate::Result<bool> {
40        if self.element_count() != other.element_count() {
41            return Ok(false);
42        }
43        D::i_allclose(&*self.storage_read()?, self.layout(), &*other.storage_read()?, other.layout())
44    }
45}
46
47impl<D: Device> Tensor<D, Bool> {
48    pub fn allclose(&self, other: &Self) -> crate::Result<bool> {
49        if self.element_count() != other.element_count() {
50            return Ok(false);
51        }
52        D::b_allclose(&*self.storage_read()?, self.layout(), &*other.storage_read()?, other.layout())
53    }
54}
55
56// ---- false_count (Bool) ----
57
58impl<D: Device> Tensor<D, Bool> {
59    pub fn false_count(&self) -> crate::Result<usize> {
60        Ok(self.element_count() - self.true_count()?)
61    }
62}
63
64// ---- Display ----
65//
66// Prints in NumPy/PyTorch style:
67//   scalar:   tensor(3.14)
68//   1-D:      tensor([1., 2., 3.])
69//   2-D:      tensor([[1., 2.],
70//                     [3., 4.]])
71//   Higher:   nested [ ... ]
72//
73// Precision: 4 significant digits for float, exact for int/bool.
74
75fn fmt_f64(v: f64, precision: usize) -> String {
76    let abs = v.abs();
77    if abs == 0.0 || (abs >= 0.001 && abs < 1e5) {
78        format!("{:.prec$}", v, prec = precision)
79    } else {
80        format!("{:.prec$e}", v, prec = precision)
81    }
82}
83
84fn write_nd(buf: &mut String, data: &[String], dims: &[usize], depth: usize, indent: usize) {
85    if dims.is_empty() {
86        // scalar
87        buf.push_str(&data[0]);
88        return;
89    }
90    if dims.len() == 1 {
91        buf.push('[');
92        for (i, v) in data.iter().enumerate() {
93            if i > 0 {
94                buf.push_str(", ");
95            }
96            buf.push_str(v);
97        }
98        buf.push(']');
99        return;
100    }
101    // multi-dim: recurse
102    let stride: usize = dims[1..].iter().product();
103    buf.push('[');
104    for i in 0..dims[0] {
105        if i > 0 {
106            buf.push(',');
107            buf.push('\n');
108            for _ in 0..indent + depth + 1 {
109                buf.push(' ');
110            }
111        }
112        let slice = &data[i * stride..(i + 1) * stride];
113        write_nd(buf, slice, &dims[1..], depth + 1, indent);
114    }
115    buf.push(']');
116}
117
118impl<D: Device> std::fmt::Display for Tensor<D, Float> {
119    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
120        const PREC: usize = 4;
121        let vals = match self.to_vec() {
122            Ok(v) => v,
123            Err(e) => return write!(f, "<Tensor error: {}>", e),
124        };
125        let strs: Vec<String> = vals.iter().map(|&v| fmt_f64(v, PREC)).collect();
126        let dtype_str = match self.dtype() {
127            FloatDType::F32 => "f32",
128            FloatDType::F64 => "f64",
129        };
130        let mut buf = format!("Tensor<{}>(", dtype_str);
131        let indent = buf.len();
132        write_nd(&mut buf, &strs, self.shape().dims(), 0, indent);
133        buf.push(')');
134        write!(f, "{}", buf)
135    }
136}
137
138impl<D: Device> std::fmt::Display for Tensor<D, Int> {
139    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
140        let vals = match self.to_vec() {
141            Ok(v) => v,
142            Err(e) => return write!(f, "<Tensor error: {}>", e),
143        };
144        let strs: Vec<String> = vals.iter().map(|v| v.to_string()).collect();
145        let dtype_str = match self.dtype() {
146            IntDType::I32 => "i32",
147            IntDType::U32 => "u32",
148            IntDType::U8 => "u8",
149        };
150        let mut buf = format!("Tensor<{}>(", dtype_str);
151        let indent = buf.len();
152        write_nd(&mut buf, &strs, self.shape().dims(), 0, indent);
153        buf.push(')');
154        write!(f, "{}", buf)
155    }
156}
157
158impl<D: Device> std::fmt::Display for Tensor<D, Bool> {
159    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
160        let vals = match self.to_vec() {
161            Ok(v) => v,
162            Err(e) => return write!(f, "<Tensor error: {}>", e),
163        };
164        let strs: Vec<String> = vals.iter().map(|v| v.to_string()).collect();
165        let mut buf = "Tensor<bool>(".to_string();
166        let indent = buf.len();
167        write_nd(&mut buf, &strs, self.shape().dims(), 0, indent);
168        buf.push(')');
169        write!(f, "{}", buf)
170    }
171}
luma_tensor/ops/display.rs

luma_tensor/ops/
display.rs
