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