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 fn bool_gather(
90 dim: usize,
91 tensor: BoolTensor<Self>,
92 indices: IntTensor<Self>,
93 ) -> BoolTensor<Self> {
94 binary_op!((tensor, bool), (indices, int), |tensor, indices| B::bool_gather(dim, tensor, indices) => Bool)
95 }
96
97 fn bool_scatter_or(
98 dim: usize,
99 tensor: BoolTensor<Self>,
100 indices: IntTensor<Self>,
101 value: BoolTensor<Self>,
102 ) -> BoolTensor<Self> {
103 multi_op!(
104 inputs[(tensor, bool), (indices, int), (value, bool)], => Bool,
105 B::bool_scatter_or(dim, tensor, indices, value)
106 )
107 }
108
109 fn bool_equal(lhs: BoolTensor<Self>, rhs: BoolTensor<Self>) -> BoolTensor<Self> {
110 binary_op!((lhs, bool), (rhs, bool), |lhs, rhs| B::bool_equal(lhs, rhs) => Bool)
111 }
112
113 fn bool_equal_elem(lhs: BoolTensor<Self>, rhs: Scalar) -> BoolTensor<Self> {
114 unary_op!(lhs, bool, |lhs| B::bool_equal_elem(lhs, rhs) => Bool)
115 }
116
117 fn bool_not(tensor: BoolTensor<Self>) -> BoolTensor<Self> {
118 unary_op!(tensor, bool, |tensor| B::bool_not(tensor) => Bool)
119 }
120
121 fn bool_and(lhs: BoolTensor<Self>, rhs: BoolTensor<Self>) -> BoolTensor<Self> {
122 binary_op!((lhs, bool), (rhs, bool), |lhs, rhs| B::bool_and(lhs, rhs) => Bool)
123 }
124
125 fn bool_or(lhs: BoolTensor<Self>, rhs: BoolTensor<Self>) -> BoolTensor<Self> {
126 binary_op!((lhs, bool), (rhs, bool), |lhs, rhs| B::bool_or(lhs, rhs) => Bool)
127 }
128
129 fn bool_swap_dims(tensor: BoolTensor<Self>, dim1: usize, dim2: usize) -> BoolTensor<Self> {
130 unary_op!(tensor, bool, |tensor| B::bool_swap_dims(tensor, dim1, dim2) => Bool)
131 }
132
133 fn bool_permute(tensor: BoolTensor<Self>, axes: &[usize]) -> BoolTensor<Self> {
134 unary_op!(tensor, bool, |tensor| B::bool_permute(tensor, axes) => Bool)
135 }
136
137 fn bool_flip(tensor: BoolTensor<Self>, axes: &[usize]) -> BoolTensor<Self> {
138 unary_op!(tensor, bool, |tensor| B::bool_flip(tensor, axes) => Bool)
139 }
140
141 fn bool_expand(tensor: BoolTensor<Self>, shape: Shape) -> BoolTensor<Self> {
142 unary_op!(tensor, bool, |tensor| B::bool_expand(tensor, shape) => Bool)
143 }
144
145 fn bool_unfold(
146 tensor: BoolTensor<Self>,
147 dim: usize,
148 size: usize,
149 step: usize,
150 ) -> BoolTensor<Self> {
151 unary_op!(tensor, bool, |tensor| B::bool_unfold(tensor, dim, size, step) => Bool)
152 }
153
154 fn bool_select(
155 tensor: BoolTensor<Self>,
156 dim: usize,
157 indices: IntTensor<Self>,
158 ) -> BoolTensor<Self> {
159 binary_op!((tensor, bool), (indices, int), |tensor, indices| B::bool_select(tensor, dim, indices) => Bool)
160 }
161
162 fn bool_select_or(
163 tensor: BoolTensor<Self>,
164 dim: usize,
165 indices: IntTensor<Self>,
166 value: BoolTensor<Self>,
167 ) -> BoolTensor<Self> {
168 multi_op!(
169 inputs[(tensor, bool), (indices, int), (value, bool)], => Bool,
170 B::bool_select_or(tensor, dim, indices, value)
171 )
172 }
173
174 fn bool_repeat_dim(tensor: BoolTensor<Self>, dim: usize, times: usize) -> BoolTensor<Self> {
175 unary_op!(tensor, bool, |tensor| B::bool_repeat_dim(tensor, dim, times) => Bool)
176 }
177
178 fn bool_cat(tensors: Vec<BoolTensor<Self>>, dim: usize) -> BoolTensor<Self> {
179 vec_op!(tensors, bool, |tensors| B::bool_cat(tensors, dim) => Bool)
180 }
181
182 fn bool_not_equal(lhs: BoolTensor<Self>, rhs: BoolTensor<Self>) -> BoolTensor<Self> {
183 binary_op!((lhs, bool), (rhs, bool), |lhs, rhs| B::bool_not_equal(lhs, rhs) => Bool)
184 }
185
186 fn bool_not_equal_elem(lhs: BoolTensor<Self>, rhs: Scalar) -> BoolTensor<Self> {
187 unary_op!(lhs, bool, |lhs| B::bool_not_equal_elem(lhs, rhs) => Bool)
188 }
189
190 fn bool_xor(lhs: BoolTensor<Self>, rhs: BoolTensor<Self>) -> BoolTensor<Self> {
191 binary_op!((lhs, bool), (rhs, bool), |lhs, rhs| B::bool_xor(lhs, rhs) => Bool)
192 }
193
194 fn bool_transpose(tensor: BoolTensor<Self>) -> BoolTensor<Self> {
195 unary_op!(tensor, bool, |tensor| B::bool_transpose(tensor) => Bool)
196 }
197
198 fn bool_any(tensor: BoolTensor<Self>) -> BoolTensor<Self> {
199 unary_op!(tensor, bool, |tensor| B::bool_any(tensor) => Bool)
200 }
201
202 fn bool_any_dim(tensor: BoolTensor<Self>, dim: usize) -> BoolTensor<Self> {
203 unary_op!(tensor, bool, |tensor| B::bool_any_dim(tensor, dim) => Bool)
204 }
205
206 fn bool_all(tensor: BoolTensor<Self>) -> BoolTensor<Self> {
207 unary_op!(tensor, bool, |tensor| B::bool_all(tensor) => Bool)
208 }
209
210 fn bool_all_dim(tensor: BoolTensor<Self>, dim: usize) -> BoolTensor<Self> {
211 unary_op!(tensor, bool, |tensor| B::bool_all_dim(tensor, dim) => Bool)
212 }
213
214 async fn bool_argwhere(tensor: BoolTensor<Self>, out_dtype: IntDType) -> IntTensor<Self> {
215 unary_op!(tensor, bool, |tensor| B::bool_argwhere(tensor, out_dtype).await => Int)
216 }
217}