burn_dispatch/ops/
bool_tensor.rs

1use alloc::vec::Vec;
2use burn_backend::{
3    BoolDType, ExecutionError, FloatDType, IntDType, Scalar, Shape, Slice, TensorData,
4    ops::BoolTensorOps,
5    tensor::{BoolTensor, FloatTensor, IndexingUpdateOp, IntTensor},
6};
7use burn_backend_extension::backend_dispatch;
8
9use crate::{Dispatch, DispatchDevice};
10
11#[backend_dispatch]
12impl BoolTensorOps<Self> for Dispatch {
13    fn bool_empty(shape: Shape, device: &DispatchDevice, dtype: BoolDType) -> BoolTensor<Self> {
14        B::bool_empty(shape, device, dtype)
15    }
16
17    fn bool_zeros(shape: Shape, device: &DispatchDevice, dtype: BoolDType) -> BoolTensor<Self> {
18        B::bool_zeros(shape, device, dtype)
19    }
20
21    fn bool_ones(shape: Shape, device: &DispatchDevice, dtype: BoolDType) -> BoolTensor<Self> {
22        B::bool_ones(shape, device, dtype)
23    }
24
25    async fn bool_into_data(tensor: BoolTensor<Self>) -> Result<TensorData, ExecutionError> {
26        B::bool_into_data(tensor).await
27    }
28
29    fn bool_from_data(data: TensorData, device: &DispatchDevice) -> BoolTensor<Self> {
30        B::bool_from_data(data, device)
31    }
32
33    fn bool_into_int(tensor: BoolTensor<Self>, out_dtype: IntDType) -> IntTensor<Self> {
34        B::bool_into_int(tensor, out_dtype)
35    }
36
37    fn bool_into_float(tensor: BoolTensor<Self>, out_dtype: FloatDType) -> FloatTensor<Self> {
38        B::bool_into_float(tensor, out_dtype)
39    }
40
41    #[backend_dispatch(skip)]
42    fn bool_to_device(tensor: BoolTensor<Self>, device: &DispatchDevice) -> BoolTensor<Self> {
43        to_device!(
44            Bool,
45            bool,
46            tensor,
47            device,
48            bool_to_device,
49            |inner, device| {
50                let data =
51                    burn_backend::read_sync(B1::bool_into_data(inner)).expect("Should read data");
52                B2::bool_from_data(data, device)
53            }
54        )
55    }
56
57    fn bool_reshape(tensor: BoolTensor<Self>, shape: Shape) -> BoolTensor<Self> {
58        B::bool_reshape(tensor, shape)
59    }
60
61    fn bool_slice(tensor: BoolTensor<Self>, slices: &[Slice]) -> BoolTensor<Self> {
62        B::bool_slice(tensor, slices)
63    }
64
65    fn bool_slice_assign(
66        tensor: BoolTensor<Self>,
67        slices: &[Slice],
68        value: BoolTensor<Self>,
69    ) -> BoolTensor<Self> {
70        B::bool_slice_assign(tensor, slices, value)
71    }
72
73    fn bool_mask_where(
74        tensor: BoolTensor<Self>,
75        mask: BoolTensor<Self>,
76        value: BoolTensor<Self>,
77    ) -> BoolTensor<Self> {
78        B::bool_mask_where(tensor, mask, value)
79    }
80
81    fn bool_mask_fill(
82        tensor: BoolTensor<Self>,
83        mask: BoolTensor<Self>,
84        value: Scalar,
85    ) -> BoolTensor<Self> {
86        B::bool_mask_fill(tensor, mask, value)
87    }
88
89    async fn bool_mask_select(
90        tensor: BoolTensor<Self>,
91        mask: BoolTensor<Self>,
92    ) -> BoolTensor<Self> {
93        B::bool_mask_select(tensor, mask).await
94    }
95
96    fn bool_gather(
97        dim: usize,
98        tensor: BoolTensor<Self>,
99        indices: IntTensor<Self>,
100    ) -> BoolTensor<Self> {
101        B::bool_gather(dim, tensor, indices)
102    }
103
104    fn bool_scatter(
105        dim: usize,
106        tensor: BoolTensor<Self>,
107        indices: IntTensor<Self>,
108        value: BoolTensor<Self>,
109        update: IndexingUpdateOp,
110    ) -> BoolTensor<Self> {
111        B::bool_scatter(dim, tensor, indices, value, update)
112    }
113
114    fn bool_scatter_nd(
115        data: BoolTensor<Self>,
116        indices: IntTensor<Self>,
117        values: BoolTensor<Self>,
118        reduction: IndexingUpdateOp,
119    ) -> BoolTensor<Self> {
120        B::bool_scatter_nd(data, indices, values, reduction)
121    }
122
123    fn bool_gather_nd(data: BoolTensor<Self>, indices: IntTensor<Self>) -> BoolTensor<Self> {
124        B::bool_gather_nd(data, indices)
125    }
126
127    fn bool_equal(lhs: BoolTensor<Self>, rhs: BoolTensor<Self>) -> BoolTensor<Self> {
128        B::bool_equal(lhs, rhs)
129    }
130
131    fn bool_equal_elem(lhs: BoolTensor<Self>, rhs: Scalar) -> BoolTensor<Self> {
132        B::bool_equal_elem(lhs, rhs)
133    }
134
135    fn bool_not(tensor: BoolTensor<Self>) -> BoolTensor<Self> {
136        B::bool_not(tensor)
137    }
138
139    fn bool_and(lhs: BoolTensor<Self>, rhs: BoolTensor<Self>) -> BoolTensor<Self> {
140        B::bool_and(lhs, rhs)
141    }
142
143    fn bool_or(lhs: BoolTensor<Self>, rhs: BoolTensor<Self>) -> BoolTensor<Self> {
144        B::bool_or(lhs, rhs)
145    }
146
147    fn bool_swap_dims(tensor: BoolTensor<Self>, dim1: usize, dim2: usize) -> BoolTensor<Self> {
148        B::bool_swap_dims(tensor, dim1, dim2)
149    }
150
151    fn bool_permute(tensor: BoolTensor<Self>, axes: &[usize]) -> BoolTensor<Self> {
152        B::bool_permute(tensor, axes)
153    }
154
155    fn bool_flip(tensor: BoolTensor<Self>, axes: &[usize]) -> BoolTensor<Self> {
156        B::bool_flip(tensor, axes)
157    }
158
159    fn bool_expand(tensor: BoolTensor<Self>, shape: Shape) -> BoolTensor<Self> {
160        B::bool_expand(tensor, shape)
161    }
162
163    fn bool_unfold(
164        tensor: BoolTensor<Self>,
165        dim: usize,
166        size: usize,
167        step: usize,
168    ) -> BoolTensor<Self> {
169        B::bool_unfold(tensor, dim, size, step)
170    }
171
172    fn bool_select(
173        tensor: BoolTensor<Self>,
174        dim: usize,
175        indices: IntTensor<Self>,
176    ) -> BoolTensor<Self> {
177        B::bool_select(tensor, dim, indices)
178    }
179
180    fn bool_select_assign(
181        tensor: BoolTensor<Self>,
182        dim: usize,
183        indices: IntTensor<Self>,
184        value: BoolTensor<Self>,
185        update: IndexingUpdateOp,
186    ) -> BoolTensor<Self> {
187        B::bool_select_assign(tensor, dim, indices, value, update)
188    }
189
190    fn bool_repeat_dim(tensor: BoolTensor<Self>, dim: usize, times: usize) -> BoolTensor<Self> {
191        B::bool_repeat_dim(tensor, dim, times)
192    }
193
194    fn bool_cat(tensors: Vec<BoolTensor<Self>>, dim: usize) -> BoolTensor<Self> {
195        B::bool_cat(tensors, dim)
196    }
197
198    fn bool_not_equal(lhs: BoolTensor<Self>, rhs: BoolTensor<Self>) -> BoolTensor<Self> {
199        B::bool_not_equal(lhs, rhs)
200    }
201
202    fn bool_not_equal_elem(lhs: BoolTensor<Self>, rhs: Scalar) -> BoolTensor<Self> {
203        B::bool_not_equal_elem(lhs, rhs)
204    }
205
206    fn bool_xor(lhs: BoolTensor<Self>, rhs: BoolTensor<Self>) -> BoolTensor<Self> {
207        B::bool_xor(lhs, rhs)
208    }
209
210    fn bool_transpose(tensor: BoolTensor<Self>) -> BoolTensor<Self> {
211        B::bool_transpose(tensor)
212    }
213
214    fn bool_any(tensor: BoolTensor<Self>) -> BoolTensor<Self> {
215        B::bool_any(tensor)
216    }
217
218    fn bool_any_dim(tensor: BoolTensor<Self>, dim: usize) -> BoolTensor<Self> {
219        B::bool_any_dim(tensor, dim)
220    }
221
222    fn bool_all(tensor: BoolTensor<Self>) -> BoolTensor<Self> {
223        B::bool_all(tensor)
224    }
225
226    fn bool_all_dim(tensor: BoolTensor<Self>, dim: usize) -> BoolTensor<Self> {
227        B::bool_all_dim(tensor, dim)
228    }
229
230    async fn bool_argwhere(tensor: BoolTensor<Self>, out_dtype: IntDType) -> IntTensor<Self> {
231        B::bool_argwhere(tensor, out_dtype).await
232    }
233}
burn_dispatch/ops/bool_tensor.rs

burn_dispatch/ops/
bool_tensor.rs
