burn_ndarray/ops/
bool_tensor.rs

1// Language
2use alloc::vec;
3use alloc::vec::Vec;
4use burn_backend::Scalar;
5use burn_backend::{ElementConversion, TensorMetadata, tensor::FloatTensor};
6use burn_backend::{
7    backend::ExecutionError,
8    ops::BoolTensorOps,
9    tensor::{BoolTensor, IntTensor},
10};
11use burn_std::{BoolDType, FloatDType, IntDType};
12use ndarray::IntoDimension;
13
14// Current crate
15use crate::{NdArray, execute_with_int_dtype, tensor::NdArrayTensor};
16use crate::{
17    NdArrayDevice, SharedArray, execute_with_float_out_dtype, execute_with_int_out_dtype, slice,
18};
19
20// Workspace crates
21use burn_backend::{Shape, TensorData};
22
23use super::{NdArrayBoolOps, NdArrayOps};
24
25impl BoolTensorOps<Self> for NdArray {
26    fn bool_from_data(data: TensorData, _device: &NdArrayDevice) -> NdArrayTensor {
27        if !data.dtype.is_bool() {
28            unimplemented!("Unsupported dtype for `bool_from_data`")
29        }
30        NdArrayTensor::from_data(data)
31    }
32
33    async fn bool_into_data(tensor: NdArrayTensor) -> Result<TensorData, ExecutionError> {
34        Ok(tensor.into_data())
35    }
36
37    fn bool_to_device(tensor: NdArrayTensor, _device: &NdArrayDevice) -> NdArrayTensor {
38        tensor
39    }
40
41    fn bool_reshape(tensor: NdArrayTensor, shape: Shape) -> NdArrayTensor {
42        NdArrayOps::reshape(tensor.bool(), shape).into()
43    }
44
45    fn bool_slice(tensor: NdArrayTensor, slices: &[burn_backend::Slice]) -> NdArrayTensor {
46        slice!(tensor, slices)
47    }
48
49    fn bool_into_int(tensor: NdArrayTensor, out_dtype: IntDType) -> NdArrayTensor {
50        // Use mapv directly instead of collecting to Vec and going through TensorData
51        execute_with_int_out_dtype!(
52            out_dtype,
53            I,
54            tensor.bool().mapv(|b| b.elem::<I>()).into_shared().into()
55        )
56    }
57
58    fn bool_empty(shape: Shape, _device: &NdArrayDevice, dtype: BoolDType) -> NdArrayTensor {
59        Self::bool_zeros(shape, _device, dtype)
60    }
61
62    fn bool_zeros(shape: Shape, _device: &NdArrayDevice, _dtype: BoolDType) -> NdArrayTensor {
63        let values = vec![false; shape.num_elements()];
64        NdArrayTensor::from_data(TensorData::new(values, shape))
65    }
66
67    fn bool_ones(shape: Shape, _device: &NdArrayDevice, _dtype: BoolDType) -> NdArrayTensor {
68        let values = vec![true; shape.num_elements()];
69        NdArrayTensor::from_data(TensorData::new(values, shape))
70    }
71
72    fn bool_slice_assign(
73        tensor: NdArrayTensor,
74        slices: &[burn_backend::Slice],
75        value: NdArrayTensor,
76    ) -> NdArrayTensor {
77        NdArrayOps::slice_assign(tensor.bool(), slices, value.bool()).into()
78    }
79
80    fn bool_cat(tensors: Vec<NdArrayTensor>, dim: usize) -> NdArrayTensor {
81        NdArrayOps::cat(tensors.into_iter().map(|it| it.bool()).collect(), dim).into()
82    }
83
84    fn bool_equal(lhs: NdArrayTensor, rhs: NdArrayTensor) -> NdArrayTensor {
85        NdArrayBoolOps::equal(lhs.bool(), rhs.bool()).into()
86    }
87
88    fn bool_not(tensor: NdArrayTensor) -> NdArrayTensor {
89        tensor.bool().mapv(|a| !a).into_shared().into()
90    }
91
92    fn bool_and(lhs: NdArrayTensor, rhs: NdArrayTensor) -> NdArrayTensor {
93        NdArrayBoolOps::and(lhs.bool(), rhs.bool()).into()
94    }
95
96    fn bool_or(lhs: NdArrayTensor, rhs: NdArrayTensor) -> NdArrayTensor {
97        NdArrayBoolOps::or(lhs.bool(), rhs.bool()).into()
98    }
99
100    fn bool_into_float(tensor: NdArrayTensor, out_dtype: FloatDType) -> FloatTensor<Self> {
101        execute_with_float_out_dtype!(
102            out_dtype,
103            E,
104            tensor.bool().mapv(|b| b.elem::<E>()).into_shared().into()
105        )
106    }
107
108    fn bool_swap_dims(tensor: NdArrayTensor, dim1: usize, dim2: usize) -> NdArrayTensor {
109        NdArrayOps::swap_dims(tensor.bool(), dim1, dim2).into()
110    }
111
112    fn bool_permute(tensor: NdArrayTensor, axes: &[usize]) -> NdArrayTensor {
113        tensor.bool().permuted_axes(axes.into_dimension()).into()
114    }
115
116    fn bool_expand(tensor: NdArrayTensor, shape: Shape) -> NdArrayTensor {
117        NdArrayOps::expand(tensor.bool(), shape).into()
118    }
119
120    fn bool_select(tensor: NdArrayTensor, dim: usize, indices: NdArrayTensor) -> NdArrayTensor {
121        execute_with_int_dtype!(indices, I, |indices: SharedArray<I>| -> NdArrayTensor {
122            let tensor_bool = tensor.bool();
123            let indices_vec: Vec<usize> = indices
124                .into_iter()
125                .map(|i| i.elem::<i64>() as usize)
126                .collect();
127
128            let selected = tensor_bool.select(ndarray::Axis(dim), &indices_vec);
129            selected.into_shared().into()
130        })
131    }
132
133    fn bool_select_or(
134        tensor: NdArrayTensor,
135        dim: usize,
136        indices: NdArrayTensor,
137        value: NdArrayTensor,
138    ) -> NdArrayTensor {
139        execute_with_int_dtype!(indices, I, |indices: SharedArray<I>| -> NdArrayTensor {
140            let mut output_array = tensor.bool().into_owned();
141            let value_bool = value.bool();
142
143            for (index_value, index) in indices.into_iter().enumerate() {
144                let index_usize = index.elem::<i64>() as usize;
145                let mut view = output_array.index_axis_mut(ndarray::Axis(dim), index_usize);
146                let value_slice = value_bool.index_axis(ndarray::Axis(dim), index_value);
147                // For boolean tensors, select_assign should use logical OR operation
148                view.zip_mut_with(&value_slice, |a, b| *a = *a || *b);
149            }
150            output_array.into_shared().into()
151        })
152    }
153
154    fn bool_flip(tensor: NdArrayTensor, axes: &[usize]) -> NdArrayTensor {
155        NdArrayOps::flip(tensor.bool(), axes).into()
156    }
157
158    fn bool_unfold(tensor: NdArrayTensor, dim: usize, size: usize, step: usize) -> NdArrayTensor {
159        NdArrayOps::unfold(tensor.bool(), dim, size, step).into()
160    }
161
162    fn bool_mask_where(
163        tensor: BoolTensor<Self>,
164        mask: BoolTensor<Self>,
165        value: BoolTensor<Self>,
166    ) -> BoolTensor<Self> {
167        NdArrayOps::mask_where(tensor.bool(), mask.bool(), value.bool()).into()
168    }
169
170    fn bool_mask_fill(
171        tensor: BoolTensor<Self>,
172        mask: BoolTensor<Self>,
173        value: Scalar,
174    ) -> BoolTensor<Self> {
175        NdArrayOps::mask_fill(tensor.bool(), mask.bool(), value.elem()).into()
176    }
177
178    fn bool_gather(
179        dim: usize,
180        tensor: BoolTensor<Self>,
181        indices: IntTensor<Self>,
182    ) -> BoolTensor<Self> {
183        execute_with_int_dtype!(indices, |indices| NdArrayOps::gather(
184            dim,
185            tensor.bool(),
186            indices
187        ))
188    }
189
190    fn bool_scatter_or(
191        dim: usize,
192        tensor: BoolTensor<Self>,
193        indices: IntTensor<Self>,
194        value: BoolTensor<Self>,
195    ) -> BoolTensor<Self> {
196        execute_with_int_dtype!(indices, |indices| NdArrayOps::scatter(
197            dim,
198            tensor.bool(),
199            indices,
200            value.bool()
201        ))
202    }
203
204    fn bool_equal_elem(lhs: BoolTensor<Self>, rhs: Scalar) -> BoolTensor<Self> {
205        NdArrayBoolOps::equal_elem(lhs.bool(), rhs.elem()).into()
206    }
207
208    fn bool_any(tensor: BoolTensor<Self>) -> BoolTensor<Self> {
209        // Use view() for zero-copy on borrowed storage with short-circuit evaluation
210        let result = NdArrayBoolOps::any_view(tensor.bool().view());
211        NdArrayTensor::from_data(TensorData::new(vec![result], Shape::new([1])))
212    }
213
214    fn bool_all(tensor: BoolTensor<Self>) -> BoolTensor<Self> {
215        // Use view() for zero-copy on borrowed storage with short-circuit evaluation
216        let result = NdArrayBoolOps::all_view(tensor.bool().view());
217        NdArrayTensor::from_data(TensorData::new(vec![result], Shape::new([1])))
218    }
219}
burn_ndarray/ops/bool_tensor.rs

burn_ndarray/ops/
bool_tensor.rs
