1use super::*;
2
3impl BoolTensorOps<Self> for Dispatch {
4 fn bool_autodiff_context(tensor: &BoolTensor<Self>) -> ruda_tensor::AutodiffContext {
5 tensor.context()
6 }
7
8 fn bool_with_autodiff_context(tensor: BoolTensor<Self>, context: ruda_tensor::AutodiffContext) -> BoolTensor<Self> {
9 tensor.with_context(context)
10 }
11
12 fn bool_empty(shape: Shape, device: &Device<Self>, dtype: BoolDType) -> BoolTensor<Self> {
13 let selection = Selection::factory(device);
14 macro_rules! call {
15 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
16 let result = <$B>::bool_empty(shape, $device, dtype);
17 pack!($mode, $variant, Bool, result)
18 }};
19 }
20 dispatch!(selection, call)
21 }
22
23 fn bool_zeros(shape: Shape, device: &Device<Self>, dtype: BoolDType) -> BoolTensor<Self> {
24 let selection = Selection::factory(device);
25 macro_rules! call {
26 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
27 let result = <$B>::bool_zeros(shape, $device, dtype);
28 pack!($mode, $variant, Bool, result)
29 }};
30 }
31 dispatch!(selection, call)
32 }
33
34 fn bool_ones(shape: Shape, device: &Device<Self>, dtype: BoolDType) -> BoolTensor<Self> {
35 let selection = Selection::factory(device);
36 macro_rules! call {
37 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
38 let result = <$B>::bool_ones(shape, $device, dtype);
39 pack!($mode, $variant, Bool, result)
40 }};
41 }
42 dispatch!(selection, call)
43 }
44
45 fn bool_into_data( tensor: BoolTensor<Self>, ) -> impl Future<Output = Result<TensorData, ExecutionError>> + Send {
46 let selection = Selection::tensors(core::iter::once(&tensor));
47 macro_rules! call {
48 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
49 let future = <$B>::bool_into_data(unpack!($mode, $variant, Bool, tensor, $context));
50 let future: core::pin::Pin<Box<dyn Future<Output = Result<TensorData, ExecutionError>> + Send>> = Box::pin(async move { future.await });
51 future
52 }};
53 }
54 dispatch!(selection, call)
55 }
56
57 fn bool_from_data(data: TensorData, device: &Device<Self>) -> BoolTensor<Self> {
58 let selection = Selection::factory(device);
59 macro_rules! call {
60 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
61 let result = <$B>::bool_from_data(data, $device);
62 pack!($mode, $variant, Bool, result)
63 }};
64 }
65 dispatch!(selection, call)
66 }
67
68 fn bool_into_int(tensor: BoolTensor<Self>, out_dtype: IntDType) -> IntTensor<Self> {
69 let selection = Selection::tensors(core::iter::once(&tensor));
70 macro_rules! call {
71 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
72 let result = <$B>::bool_into_int(unpack!($mode, $variant, Bool, tensor, $context), out_dtype);
73 pack!($mode, $variant, Int, result)
74 }};
75 }
76 dispatch!(selection, call)
77 }
78
79 fn bool_into_float(tensor: BoolTensor<Self>, out_dtype: FloatDType) -> FloatTensor<Self> {
80 let selection = Selection::tensors(core::iter::once(&tensor));
81 macro_rules! call {
82 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
83 let result = <$B>::bool_into_float(unpack!($mode, $variant, Bool, tensor, $context), out_dtype);
84 pack!($mode, $variant, Float, result)
85 }};
86 }
87 dispatch!(selection, call)
88 }
89
90 fn bool_device(tensor: &BoolTensor<Self>) -> Device<Self> {
91 tensor.device()
92 }
93
94 fn bool_to_device(tensor: BoolTensor<Self>, device: &Device<Self>) -> BoolTensor<Self> {
95 crate::transfer::bool_to_device(tensor, device)
96 }
97
98 fn bool_reshape(tensor: BoolTensor<Self>, shape: Shape) -> BoolTensor<Self> {
99 let selection = Selection::tensors(core::iter::once(&tensor));
100 macro_rules! call {
101 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
102 let result = <$B>::bool_reshape(unpack!($mode, $variant, Bool, tensor, $context), shape);
103 pack!($mode, $variant, Bool, result)
104 }};
105 }
106 dispatch!(selection, call)
107 }
108
109 fn bool_slice(tensor: BoolTensor<Self>, slices: &[Slice]) -> BoolTensor<Self> {
110 let selection = Selection::tensors(core::iter::once(&tensor));
111 macro_rules! call {
112 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
113 let result = <$B>::bool_slice(unpack!($mode, $variant, Bool, tensor, $context), slices);
114 pack!($mode, $variant, Bool, result)
115 }};
116 }
117 dispatch!(selection, call)
118 }
119
120 fn bool_slice_assign( tensor: BoolTensor<Self>, slices: &[Slice], value: BoolTensor<Self>, ) -> BoolTensor<Self> {
121 let selection = Selection::tensors(core::iter::once(&tensor).chain(core::iter::once(&value)));
122 macro_rules! call {
123 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
124 let result = <$B>::bool_slice_assign(unpack!($mode, $variant, Bool, tensor, $context), slices, unpack!($mode, $variant, Bool, value, $context));
125 pack!($mode, $variant, Bool, result)
126 }};
127 }
128 dispatch!(selection, call)
129 }
130
131 fn bool_mask_where( tensor: BoolTensor<Self>, mask: BoolTensor<Self>, value: BoolTensor<Self>, ) -> BoolTensor<Self> {
132 let selection = Selection::tensors(core::iter::once(&tensor).chain(core::iter::once(&mask)).chain(core::iter::once(&value)));
133 macro_rules! call {
134 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
135 let result = <$B>::bool_mask_where(unpack!($mode, $variant, Bool, tensor, $context), unpack!($mode, $variant, Bool, mask, $context), unpack!($mode, $variant, Bool, value, $context));
136 pack!($mode, $variant, Bool, result)
137 }};
138 }
139 dispatch!(selection, call)
140 }
141
142 fn bool_mask_fill(tensor: BoolTensor<Self>, mask: BoolTensor<Self>, value: Scalar) -> BoolTensor<Self> {
143 let selection = Selection::tensors(core::iter::once(&tensor).chain(core::iter::once(&mask)));
144 macro_rules! call {
145 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
146 let result = <$B>::bool_mask_fill(unpack!($mode, $variant, Bool, tensor, $context), unpack!($mode, $variant, Bool, mask, $context), value);
147 pack!($mode, $variant, Bool, result)
148 }};
149 }
150 dispatch!(selection, call)
151 }
152
153 fn bool_gather(dim: usize, tensor: BoolTensor<Self>, indices: IntTensor<Self>) -> BoolTensor<Self> {
154 let selection = Selection::tensors(core::iter::once(&tensor).chain(core::iter::once(&indices)));
155 macro_rules! call {
156 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
157 let result = <$B>::bool_gather(dim, unpack!($mode, $variant, Bool, tensor, $context), unpack!($mode, $variant, Int, indices, $context));
158 pack!($mode, $variant, Bool, result)
159 }};
160 }
161 dispatch!(selection, call)
162 }
163
164 fn bool_scatter_or( dim: usize, tensor: BoolTensor<Self>, indices: IntTensor<Self>, value: BoolTensor<Self>, ) -> BoolTensor<Self> {
165 let selection = Selection::tensors(core::iter::once(&tensor).chain(core::iter::once(&indices)).chain(core::iter::once(&value)));
166 macro_rules! call {
167 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
168 let result = <$B>::bool_scatter_or(dim, unpack!($mode, $variant, Bool, tensor, $context), unpack!($mode, $variant, Int, indices, $context), unpack!($mode, $variant, Bool, value, $context));
169 pack!($mode, $variant, Bool, result)
170 }};
171 }
172 dispatch!(selection, call)
173 }
174
175 fn bool_select(tensor: BoolTensor<Self>, dim: usize, indices: IntTensor<Self>) -> BoolTensor<Self> {
176 let selection = Selection::tensors(core::iter::once(&tensor).chain(core::iter::once(&indices)));
177 macro_rules! call {
178 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
179 let result = <$B>::bool_select(unpack!($mode, $variant, Bool, tensor, $context), dim, unpack!($mode, $variant, Int, indices, $context));
180 pack!($mode, $variant, Bool, result)
181 }};
182 }
183 dispatch!(selection, call)
184 }
185
186 fn bool_select_or( tensor: BoolTensor<Self>, dim: usize, indices: IntTensor<Self>, value: BoolTensor<Self>, ) -> BoolTensor<Self> {
187 let selection = Selection::tensors(core::iter::once(&tensor).chain(core::iter::once(&indices)).chain(core::iter::once(&value)));
188 macro_rules! call {
189 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
190 let result = <$B>::bool_select_or(unpack!($mode, $variant, Bool, tensor, $context), dim, unpack!($mode, $variant, Int, indices, $context), unpack!($mode, $variant, Bool, value, $context));
191 pack!($mode, $variant, Bool, result)
192 }};
193 }
194 dispatch!(selection, call)
195 }
196
197 fn bool_repeat_dim(tensor: BoolTensor<Self>, dim: usize, times: usize) -> BoolTensor<Self> {
198 let selection = Selection::tensors(core::iter::once(&tensor));
199 macro_rules! call {
200 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
201 let result = <$B>::bool_repeat_dim(unpack!($mode, $variant, Bool, tensor, $context), dim, times);
202 pack!($mode, $variant, Bool, result)
203 }};
204 }
205 dispatch!(selection, call)
206 }
207
208 fn bool_cat(tensors: Vec<BoolTensor<Self>>, dim: usize) -> BoolTensor<Self> {
209 let selection = Selection::tensors(tensors.iter());
210 macro_rules! call {
211 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
212 let result = <$B>::bool_cat(tensors.into_iter().map(|value| unpack!($mode, $variant, Bool, value, $context)).collect(), dim);
213 pack!($mode, $variant, Bool, result)
214 }};
215 }
216 dispatch!(selection, call)
217 }
218
219 fn bool_equal(lhs: BoolTensor<Self>, rhs: BoolTensor<Self>) -> BoolTensor<Self> {
220 let selection = Selection::tensors(core::iter::once(&lhs).chain(core::iter::once(&rhs)));
221 macro_rules! call {
222 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
223 let result = <$B>::bool_equal(unpack!($mode, $variant, Bool, lhs, $context), unpack!($mode, $variant, Bool, rhs, $context));
224 pack!($mode, $variant, Bool, result)
225 }};
226 }
227 dispatch!(selection, call)
228 }
229
230 fn bool_not_equal(lhs: BoolTensor<Self>, rhs: BoolTensor<Self>) -> BoolTensor<Self> {
231 let selection = Selection::tensors(core::iter::once(&lhs).chain(core::iter::once(&rhs)));
232 macro_rules! call {
233 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
234 let result = <$B>::bool_not_equal(unpack!($mode, $variant, Bool, lhs, $context), unpack!($mode, $variant, Bool, rhs, $context));
235 pack!($mode, $variant, Bool, result)
236 }};
237 }
238 dispatch!(selection, call)
239 }
240
241 fn bool_equal_elem(lhs: BoolTensor<Self>, rhs: Scalar) -> BoolTensor<Self> {
242 let selection = Selection::tensors(core::iter::once(&lhs));
243 macro_rules! call {
244 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
245 let result = <$B>::bool_equal_elem(unpack!($mode, $variant, Bool, lhs, $context), rhs);
246 pack!($mode, $variant, Bool, result)
247 }};
248 }
249 dispatch!(selection, call)
250 }
251
252 fn bool_not_equal_elem(lhs: BoolTensor<Self>, rhs: Scalar) -> BoolTensor<Self> {
253 let selection = Selection::tensors(core::iter::once(&lhs));
254 macro_rules! call {
255 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
256 let result = <$B>::bool_not_equal_elem(unpack!($mode, $variant, Bool, lhs, $context), rhs);
257 pack!($mode, $variant, Bool, result)
258 }};
259 }
260 dispatch!(selection, call)
261 }
262
263 fn bool_not(tensor: BoolTensor<Self>) -> BoolTensor<Self> {
264 let selection = Selection::tensors(core::iter::once(&tensor));
265 macro_rules! call {
266 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
267 let result = <$B>::bool_not(unpack!($mode, $variant, Bool, tensor, $context));
268 pack!($mode, $variant, Bool, result)
269 }};
270 }
271 dispatch!(selection, call)
272 }
273
274 fn bool_and(lhs: BoolTensor<Self>, rhs: BoolTensor<Self>) -> BoolTensor<Self> {
275 let selection = Selection::tensors(core::iter::once(&lhs).chain(core::iter::once(&rhs)));
276 macro_rules! call {
277 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
278 let result = <$B>::bool_and(unpack!($mode, $variant, Bool, lhs, $context), unpack!($mode, $variant, Bool, rhs, $context));
279 pack!($mode, $variant, Bool, result)
280 }};
281 }
282 dispatch!(selection, call)
283 }
284
285 fn bool_or(lhs: BoolTensor<Self>, rhs: BoolTensor<Self>) -> BoolTensor<Self> {
286 let selection = Selection::tensors(core::iter::once(&lhs).chain(core::iter::once(&rhs)));
287 macro_rules! call {
288 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
289 let result = <$B>::bool_or(unpack!($mode, $variant, Bool, lhs, $context), unpack!($mode, $variant, Bool, rhs, $context));
290 pack!($mode, $variant, Bool, result)
291 }};
292 }
293 dispatch!(selection, call)
294 }
295
296 fn bool_xor(lhs: BoolTensor<Self>, rhs: BoolTensor<Self>) -> BoolTensor<Self> {
297 let selection = Selection::tensors(core::iter::once(&lhs).chain(core::iter::once(&rhs)));
298 macro_rules! call {
299 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
300 let result = <$B>::bool_xor(unpack!($mode, $variant, Bool, lhs, $context), unpack!($mode, $variant, Bool, rhs, $context));
301 pack!($mode, $variant, Bool, result)
302 }};
303 }
304 dispatch!(selection, call)
305 }
306
307 fn bool_transpose(tensor: BoolTensor<Self>) -> BoolTensor<Self> {
308 let selection = Selection::tensors(core::iter::once(&tensor));
309 macro_rules! call {
310 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
311 let result = <$B>::bool_transpose(unpack!($mode, $variant, Bool, tensor, $context));
312 pack!($mode, $variant, Bool, result)
313 }};
314 }
315 dispatch!(selection, call)
316 }
317
318 fn bool_swap_dims(tensor: BoolTensor<Self>, dim1: usize, dim2: usize) -> BoolTensor<Self> {
319 let selection = Selection::tensors(core::iter::once(&tensor));
320 macro_rules! call {
321 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
322 let result = <$B>::bool_swap_dims(unpack!($mode, $variant, Bool, tensor, $context), dim1, dim2);
323 pack!($mode, $variant, Bool, result)
324 }};
325 }
326 dispatch!(selection, call)
327 }
328
329 fn bool_permute(tensor: BoolTensor<Self>, axes: &[usize]) -> BoolTensor<Self> {
330 let selection = Selection::tensors(core::iter::once(&tensor));
331 macro_rules! call {
332 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
333 let result = <$B>::bool_permute(unpack!($mode, $variant, Bool, tensor, $context), axes);
334 pack!($mode, $variant, Bool, result)
335 }};
336 }
337 dispatch!(selection, call)
338 }
339
340 fn bool_flip(tensor: BoolTensor<Self>, axes: &[usize]) -> BoolTensor<Self> {
341 let selection = Selection::tensors(core::iter::once(&tensor));
342 macro_rules! call {
343 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
344 let result = <$B>::bool_flip(unpack!($mode, $variant, Bool, tensor, $context), axes);
345 pack!($mode, $variant, Bool, result)
346 }};
347 }
348 dispatch!(selection, call)
349 }
350
351 fn bool_any(tensor: BoolTensor<Self>) -> BoolTensor<Self> {
352 let selection = Selection::tensors(core::iter::once(&tensor));
353 macro_rules! call {
354 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
355 let result = <$B>::bool_any(unpack!($mode, $variant, Bool, tensor, $context));
356 pack!($mode, $variant, Bool, result)
357 }};
358 }
359 dispatch!(selection, call)
360 }
361
362 fn bool_any_dim(tensor: BoolTensor<Self>, dim: usize) -> BoolTensor<Self> {
363 let selection = Selection::tensors(core::iter::once(&tensor));
364 macro_rules! call {
365 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
366 let result = <$B>::bool_any_dim(unpack!($mode, $variant, Bool, tensor, $context), dim);
367 pack!($mode, $variant, Bool, result)
368 }};
369 }
370 dispatch!(selection, call)
371 }
372
373 fn bool_all(tensor: BoolTensor<Self>) -> BoolTensor<Self> {
374 let selection = Selection::tensors(core::iter::once(&tensor));
375 macro_rules! call {
376 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
377 let result = <$B>::bool_all(unpack!($mode, $variant, Bool, tensor, $context));
378 pack!($mode, $variant, Bool, result)
379 }};
380 }
381 dispatch!(selection, call)
382 }
383
384 fn bool_all_dim(tensor: BoolTensor<Self>, dim: usize) -> BoolTensor<Self> {
385 let selection = Selection::tensors(core::iter::once(&tensor));
386 macro_rules! call {
387 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
388 let result = <$B>::bool_all_dim(unpack!($mode, $variant, Bool, tensor, $context), dim);
389 pack!($mode, $variant, Bool, result)
390 }};
391 }
392 dispatch!(selection, call)
393 }
394
395 fn bool_argwhere( tensor: BoolTensor<Self>, out_dtype: IntDType, ) -> impl Future<Output = IntTensor<Self>> + 'static + Send {
396 let selection = Selection::tensors(core::iter::once(&tensor));
397 macro_rules! call {
398 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
399 let future = <$B>::bool_argwhere(unpack!($mode, $variant, Bool, tensor, $context), out_dtype);
400 let future: core::pin::Pin<Box<dyn Future<Output = IntTensor<Self>> + Send>> = Box::pin(async move { pack!($mode, $variant, Int, future.await) });
401 future
402 }};
403 }
404 dispatch!(selection, call)
405 }
406
407 fn bool_expand(tensor: BoolTensor<Self>, shape: Shape) -> BoolTensor<Self> {
408 let selection = Selection::tensors(core::iter::once(&tensor));
409 macro_rules! call {
410 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
411 let result = <$B>::bool_expand(unpack!($mode, $variant, Bool, tensor, $context), shape);
412 pack!($mode, $variant, Bool, result)
413 }};
414 }
415 dispatch!(selection, call)
416 }
417
418 fn bool_unfold(tensor: BoolTensor<Self>, dim: usize, size: usize, step: usize) -> BoolTensor<Self> {
419 let selection = Selection::tensors(core::iter::once(&tensor));
420 macro_rules! call {
421 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
422 let result = <$B>::bool_unfold(unpack!($mode, $variant, Bool, tensor, $context), dim, size, step);
423 pack!($mode, $variant, Bool, result)
424 }};
425 }
426 dispatch!(selection, call)
427 }
428}