Skip to main content

luma_tensor/ops/construct/
bool.rs

1use std::borrow::Cow;
2
3use crate::{Bool, Device, Result, Shape, Tensor};
4
5use super::helpers;
6use super::into_tensor::IntoTensor;
7use super::options::TensorCreationOptions;
8
9impl<D: Device> Tensor<D, Bool> {
10    pub fn new(data: impl IntoTensor<D, Bool>, device: &D) -> Result<Self> {
11        let shape = data.shape()?;
12        let storage = data.into_storage(device)?;
13        Ok(Self::from_storage(storage, shape, ()))
14    }
15
16    pub fn from_vec_bool<'a, S: Into<Shape>>(data: impl Into<Cow<'a, [bool]>>, shape: S, device: &D) -> Result<Self> {
17        let shape = shape.into();
18        let storage = D::b_from_bool(data, device)?;
19        Ok(Self::from_storage(storage, shape, ()))
20    }
21
22    pub fn falses<S: Into<Shape>>(shape: S, options: impl Into<TensorCreationOptions<D, Bool>>) -> Result<Self> {
23        let options: TensorCreationOptions<D, Bool> = options.into();
24        let shape = shape.into();
25        let storage = D::b_falses(&shape, &options.device, options.dtype)?;
26        Ok(Self::from_storage(storage, shape, ()))
27    }
28
29    pub fn trues<S: Into<Shape>>(shape: S, options: impl Into<TensorCreationOptions<D, Bool>>) -> Result<Self> {
30        let options: TensorCreationOptions<D, Bool> = options.into();
31        let shape = shape.into();
32        let storage = D::b_trues(&shape, &options.device, options.dtype)?;
33        Ok(Self::from_storage(storage, shape, ()))
34    }
35
36    pub fn from_slice<S: Into<Shape>>(data: &[bool], shape: S, options: impl Into<TensorCreationOptions<D, Bool>>) -> Result<Self> {
37        let options: TensorCreationOptions<D, Bool> = options.into();
38        let shape = shape.into();
39        if shape.element_count() != data.len() {
40            return Err(crate::Error::ElementSizeMismatch { expected: data.len(), got: shape.element_count(), op: "from_slice" });
41        }
42        let storage = D::b_from_bool(data, &options.device)?;
43        Ok(Self::from_storage(storage, shape, ()))
44    }
45
46    pub fn to_vec(&self) -> crate::Result<Vec<bool>> {
47        D::b_to_vec(&*self.storage_read()?, self.layout())
48    }
49
50    pub fn eye(n: usize, options: impl Into<TensorCreationOptions<D, Bool>>) -> Result<Self> {
51        let options = options.into();
52        Self::new(helpers::fill_eye::<bool>(n), &options.device)
53    }
54
55    pub fn diag(diag: &[bool], options: impl Into<TensorCreationOptions<D, Bool>>) -> Result<Self> {
56        let options = options.into();
57        let n = diag.len();
58        let mut v = vec![false; n * n];
59        for i in 0..n {
60            v[i * n + i] = diag[i];
61        }
62        Self::new(v, &options.device)
63    }
64
65    pub fn tril(n: usize, diagonal: bool, options: impl Into<TensorCreationOptions<D, Bool>>) -> Result<Self> {
66        let options = options.into();
67        Self::new(helpers::fill_tril::<bool>(n, diagonal), &options.device)
68    }
69
70    pub fn triu(n: usize, diagonal: bool, options: impl Into<TensorCreationOptions<D, Bool>>) -> Result<Self> {
71        let options = options.into();
72        Self::new(helpers::fill_triu::<bool>(n, diagonal), &options.device)
73    }
74}