luma_tensor/ops/
boolean.rs

1//! Logical ops, reductions, and masked selection for `Bool` tensors. No autograd.
2
3use crate::{Bool, DTypeKind, Device, Float, Int, Layout, Tensor, TensorMeta};
4
5impl<D: Device> Tensor<D, Bool> {
6    pub fn and(&self, rhs: &Self) -> crate::Result<Self> {
7        let shape = self.same_shape(rhs, "and")?.clone();
8        let storage = D::b_and(&*self.storage_read()?, self.layout(), &*rhs.storage_read()?, rhs.layout())?;
9        Ok(Self::from_storage(storage, shape, ()))
10    }
11
12    pub fn or(&self, rhs: &Self) -> crate::Result<Self> {
13        let shape = self.same_shape(rhs, "or")?.clone();
14        let storage = D::b_or(&*self.storage_read()?, self.layout(), &*rhs.storage_read()?, rhs.layout())?;
15        Ok(Self::from_storage(storage, shape, ()))
16    }
17
18    pub fn xor(&self, rhs: &Self) -> crate::Result<Self> {
19        let shape = self.same_shape(rhs, "xor")?.clone();
20        let storage = D::b_xor(&*self.storage_read()?, self.layout(), &*rhs.storage_read()?, rhs.layout())?;
21        Ok(Self::from_storage(storage, shape, ()))
22    }
23
24    pub fn not(&self) -> crate::Result<Self> {
25        let storage = D::b_not(&*self.storage_read()?, self.layout())?;
26        Ok(Self::from_storage(storage, self.shape().clone(), ()))
27    }
28
29    pub fn all_all(&self) -> crate::Result<bool> {
30        Ok(self.true_count()? == self.element_count())
31    }
32
33    pub fn any_all(&self) -> crate::Result<bool> {
34        Ok(self.true_count()? > 0)
35    }
36
37    pub fn true_count(&self) -> crate::Result<usize> {
38        D::b_true_count(&*self.storage_read()?, self.layout())
39    }
40}
41
42pub trait PickDTypeKind<D: Device>: DTypeKind<D> {
43    fn pick_dispatch(
44        mask: &D::BoolStorage,
45        mask_l: &Layout,
46        on_true: &Self::Storage,
47        true_l: &Layout,
48        on_false: &Self::Storage,
49        false_l: &Layout,
50    ) -> crate::Result<Self::Storage>;
51
52    fn pick_true_dispatch(
53        mask: &D::BoolStorage,
54        mask_l: &Layout,
55        value: Self::Scalar,
56        on_false: &Self::Storage,
57        false_l: &Layout,
58    ) -> crate::Result<Self::Storage>;
59
60    fn pick_false_dispatch(
61        mask: &D::BoolStorage,
62        mask_l: &Layout,
63        on_true: &Self::Storage,
64        true_l: &Layout,
65        value: Self::Scalar,
66    ) -> crate::Result<Self::Storage>;
67}
68
69impl<D: Device> PickDTypeKind<D> for Float {
70    fn pick_dispatch(
71        mask: &<D as Device>::BoolStorage,
72        mask_l: &Layout,
73        on_true: &Self::Storage,
74        true_l: &Layout,
75        on_false: &Self::Storage,
76        false_l: &Layout,
77    ) -> crate::Result<Self::Storage> {
78        D::f_pick(mask, mask_l, on_true, true_l, on_false, false_l)
79    }
80
81    fn pick_true_dispatch(
82        mask: &<D as Device>::BoolStorage,
83        mask_l: &Layout,
84        value: Self::Scalar,
85        on_false: &Self::Storage,
86        false_l: &Layout,
87    ) -> crate::Result<Self::Storage> {
88        D::f_pick_true(mask, mask_l, value, on_false, false_l)
89    }
90
91    fn pick_false_dispatch(
92        mask: &<D as Device>::BoolStorage,
93        mask_l: &Layout,
94        on_true: &Self::Storage,
95        true_l: &Layout,
96        value: Self::Scalar,
97    ) -> crate::Result<Self::Storage> {
98        D::f_pick_false(mask, mask_l, on_true, true_l, value)
99    }
100}
101
102impl<D: Device> PickDTypeKind<D> for Int {
103    fn pick_dispatch(
104        mask: &<D as Device>::BoolStorage,
105        mask_l: &Layout,
106        on_true: &Self::Storage,
107        true_l: &Layout,
108        on_false: &Self::Storage,
109        false_l: &Layout,
110    ) -> crate::Result<Self::Storage> {
111        D::i_pick(mask, mask_l, on_true, true_l, on_false, false_l)
112    }
113
114    fn pick_true_dispatch(
115        mask: &<D as Device>::BoolStorage,
116        mask_l: &Layout,
117        value: Self::Scalar,
118        on_false: &Self::Storage,
119        false_l: &Layout,
120    ) -> crate::Result<Self::Storage> {
121        D::i_pick_true(mask, mask_l, value, on_false, false_l)
122    }
123
124    fn pick_false_dispatch(
125        mask: &<D as Device>::BoolStorage,
126        mask_l: &Layout,
127        on_true: &Self::Storage,
128        true_l: &Layout,
129        value: Self::Scalar,
130    ) -> crate::Result<Self::Storage> {
131        D::i_pick_false(mask, mask_l, on_true, true_l, value)
132    }
133}
134
135impl<D: Device> PickDTypeKind<D> for Bool {
136    fn pick_dispatch(
137        mask: &<D as Device>::BoolStorage,
138        mask_l: &Layout,
139        on_true: &Self::Storage,
140        true_l: &Layout,
141        on_false: &Self::Storage,
142        false_l: &Layout,
143    ) -> crate::Result<Self::Storage> {
144        D::b_pick(mask, mask_l, on_true, true_l, on_false, false_l)
145    }
146
147    fn pick_true_dispatch(
148        mask: &<D as Device>::BoolStorage,
149        mask_l: &Layout,
150        value: Self::Scalar,
151        on_false: &Self::Storage,
152        false_l: &Layout,
153    ) -> crate::Result<Self::Storage> {
154        D::b_pick_true(mask, mask_l, value, on_false, false_l)
155    }
156
157    fn pick_false_dispatch(
158        mask: &<D as Device>::BoolStorage,
159        mask_l: &Layout,
160        on_true: &Self::Storage,
161        true_l: &Layout,
162        value: Self::Scalar,
163    ) -> crate::Result<Self::Storage> {
164        D::b_pick_false(mask, mask_l, on_true, true_l, value)
165    }
166}
167
168impl<D: Device> Tensor<D, Bool> {
169    pub fn pick<K: PickDTypeKind<D>>(&self, on_true: &Tensor<D, K>, on_false: &Tensor<D, K>) -> crate::Result<Tensor<D, K>> {
170        let storage = K::pick_dispatch(
171            &*self.storage_read()?,
172            self.layout(),
173            &*on_true.storage_read()?,
174            on_true.layout(),
175            &*on_false.storage_read()?,
176            on_false.layout(),
177        )?;
178        let meta = K::Meta::on_pick(self, Some(on_true), Some(on_false));
179        Ok(Tensor::from_storage(storage, on_true.shape().clone(), meta))
180    }
181
182    /// `mask ? value : on_false` with a scalar true-value.
183    pub fn pick_true<K: PickDTypeKind<D>>(&self, value: K::Scalar, on_false: &Tensor<D, K>) -> crate::Result<Tensor<D, K>> {
184        let storage = K::pick_true_dispatch(&*self.storage_read()?, self.layout(), value, &*on_false.storage_read()?, on_false.layout())?;
185        let meta = K::Meta::on_pick(self, None, Some(on_false));
186        Ok(Tensor::from_storage(storage, on_false.shape().clone(), meta))
187    }
188
189    /// `mask ? on_true : value` with a scalar false-value.
190    pub fn pick_false<K: PickDTypeKind<D>>(&self, on_true: &Tensor<D, K>, value: K::Scalar) -> crate::Result<Tensor<D, K>> {
191        let storage = K::pick_false_dispatch(&*self.storage_read()?, self.layout(), &*on_true.storage_read()?, on_true.layout(), value)?;
192        let meta = K::Meta::on_pick(self, Some(on_true), None);
193        Ok(Tensor::from_storage(storage, on_true.shape().clone(), meta))
194    }
195}
luma_tensor/ops/boolean.rs

luma_tensor/ops/
boolean.rs
