1use alloc::vec::Vec;
2use burn_backend::{
3 BoolDType, ExecutionError, FloatDType, IntDType, Scalar, Shape, Slice, TensorData,
4 ops::{IntTensorOps, PadMode},
5 tensor::{BoolTensor, FloatTensor, IndexingUpdateOp, IntTensor},
6};
7use burn_backend_extension::backend_dispatch;
8
9use crate::{Dispatch, DispatchDevice};
10
11#[backend_dispatch]
12impl IntTensorOps<Self> for Dispatch {
13 fn int_pad(
14 tensor: IntTensor<Self>,
15 padding: &[(usize, usize)],
16 mode: PadMode,
17 ) -> IntTensor<Self> {
18 B::int_pad(tensor, padding, mode)
19 }
20
21 fn int_empty(shape: Shape, device: &DispatchDevice, dtype: IntDType) -> IntTensor<Self> {
22 B::int_empty(shape, device, dtype)
23 }
24
25 async fn int_into_data(tensor: IntTensor<Self>) -> Result<TensorData, ExecutionError> {
26 B::int_into_data(tensor).await
27 }
28
29 fn int_from_data(data: TensorData, device: &DispatchDevice) -> IntTensor<Self> {
30 B::int_from_data(data, device)
31 }
32
33 #[backend_dispatch(skip)]
34 fn int_to_device(tensor: IntTensor<Self>, device: &DispatchDevice) -> IntTensor<Self> {
35 to_device!(Int, int, tensor, device, int_to_device, |inner, device| {
36 let data = burn_backend::read_sync(B1::int_into_data(inner)).expect("Should read data");
37 B2::int_from_data(data, device)
38 })
39 }
40
41 fn int_reshape(tensor: IntTensor<Self>, shape: Shape) -> IntTensor<Self> {
42 B::int_reshape(tensor, shape)
43 }
44
45 fn int_slice(tensor: IntTensor<Self>, slices: &[Slice]) -> IntTensor<Self> {
46 B::int_slice(tensor, slices)
47 }
48
49 fn int_slice_assign(
50 tensor: IntTensor<Self>,
51 slices: &[Slice],
52 value: IntTensor<Self>,
53 ) -> IntTensor<Self> {
54 B::int_slice_assign(tensor, slices, value)
55 }
56
57 fn int_into_float(tensor: IntTensor<Self>, out_dtype: FloatDType) -> FloatTensor<Self> {
58 B::int_into_float(tensor, out_dtype)
59 }
60
61 fn int_mask_where(
62 tensor: IntTensor<Self>,
63 mask: BoolTensor<Self>,
64 value: IntTensor<Self>,
65 ) -> IntTensor<Self> {
66 B::int_mask_where(tensor, mask, value)
67 }
68
69 fn int_mask_fill(
70 tensor: IntTensor<Self>,
71 mask: BoolTensor<Self>,
72 value: Scalar,
73 ) -> IntTensor<Self> {
74 B::int_mask_fill(tensor, mask, value)
75 }
76
77 async fn int_mask_select(tensor: IntTensor<Self>, mask: BoolTensor<Self>) -> IntTensor<Self> {
78 B::int_mask_select(tensor, mask).await
79 }
80
81 fn int_gather(
82 dim: usize,
83 tensor: IntTensor<Self>,
84 indices: IntTensor<Self>,
85 ) -> IntTensor<Self> {
86 B::int_gather(dim, tensor, indices)
87 }
88
89 fn int_scatter(
90 dim: usize,
91 tensor: IntTensor<Self>,
92 indices: IntTensor<Self>,
93 value: IntTensor<Self>,
94 update: IndexingUpdateOp,
95 ) -> IntTensor<Self> {
96 B::int_scatter(dim, tensor, indices, value, update)
97 }
98
99 fn int_scatter_nd(
100 data: IntTensor<Self>,
101 indices: IntTensor<Self>,
102 values: IntTensor<Self>,
103 reduction: burn_backend::tensor::IndexingUpdateOp,
104 ) -> IntTensor<Self> {
105 B::int_scatter_nd(data, indices, values, reduction)
106 }
107
108 fn int_gather_nd(data: IntTensor<Self>, indices: IntTensor<Self>) -> IntTensor<Self> {
109 B::int_gather_nd(data, indices)
110 }
111
112 fn int_select(
113 tensor: IntTensor<Self>,
114 dim: usize,
115 indices: IntTensor<Self>,
116 ) -> IntTensor<Self> {
117 B::int_select(tensor, dim, indices)
118 }
119
120 fn int_select_assign(
121 tensor: IntTensor<Self>,
122 dim: usize,
123 indices: IntTensor<Self>,
124 value: IntTensor<Self>,
125 update: IndexingUpdateOp,
126 ) -> IntTensor<Self> {
127 B::int_select_assign(tensor, dim, indices, value, update)
128 }
129
130 fn int_equal(
131 lhs: IntTensor<Self>,
132 rhs: IntTensor<Self>,
133 out_dtype: BoolDType,
134 ) -> BoolTensor<Self> {
135 B::int_equal(lhs, rhs, out_dtype)
136 }
137
138 fn int_equal_elem(lhs: IntTensor<Self>, rhs: Scalar, out_dtype: BoolDType) -> BoolTensor<Self> {
139 B::int_equal_elem(lhs, rhs, out_dtype)
140 }
141
142 fn int_greater(
143 lhs: IntTensor<Self>,
144 rhs: IntTensor<Self>,
145 out_dtype: BoolDType,
146 ) -> BoolTensor<Self> {
147 B::int_greater(lhs, rhs, out_dtype)
148 }
149
150 fn int_greater_elem(
151 lhs: IntTensor<Self>,
152 rhs: Scalar,
153 out_dtype: BoolDType,
154 ) -> BoolTensor<Self> {
155 B::int_greater_elem(lhs, rhs, out_dtype)
156 }
157
158 fn int_greater_equal(
159 lhs: IntTensor<Self>,
160 rhs: IntTensor<Self>,
161 out_dtype: BoolDType,
162 ) -> BoolTensor<Self> {
163 B::int_greater_equal(lhs, rhs, out_dtype)
164 }
165
166 fn int_greater_equal_elem(
167 lhs: IntTensor<Self>,
168 rhs: Scalar,
169 out_dtype: BoolDType,
170 ) -> BoolTensor<Self> {
171 B::int_greater_equal_elem(lhs, rhs, out_dtype)
172 }
173
174 fn int_lower(
175 lhs: IntTensor<Self>,
176 rhs: IntTensor<Self>,
177 out_dtype: BoolDType,
178 ) -> BoolTensor<Self> {
179 B::int_lower(lhs, rhs, out_dtype)
180 }
181
182 fn int_lower_elem(lhs: IntTensor<Self>, rhs: Scalar, out_dtype: BoolDType) -> BoolTensor<Self> {
183 B::int_lower_elem(lhs, rhs, out_dtype)
184 }
185
186 fn int_lower_equal(
187 lhs: IntTensor<Self>,
188 rhs: IntTensor<Self>,
189 out_dtype: BoolDType,
190 ) -> BoolTensor<Self> {
191 B::int_lower_equal(lhs, rhs, out_dtype)
192 }
193
194 fn int_lower_equal_elem(
195 lhs: IntTensor<Self>,
196 rhs: Scalar,
197 out_dtype: BoolDType,
198 ) -> BoolTensor<Self> {
199 B::int_lower_equal_elem(lhs, rhs, out_dtype)
200 }
201
202 fn int_add(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
203 B::int_add(lhs, rhs)
204 }
205
206 fn int_add_scalar(lhs: IntTensor<Self>, rhs: Scalar) -> IntTensor<Self> {
207 B::int_add_scalar(lhs, rhs)
208 }
209
210 fn int_sub(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
211 B::int_sub(lhs, rhs)
212 }
213
214 fn int_sub_scalar(lhs: IntTensor<Self>, rhs: Scalar) -> IntTensor<Self> {
215 B::int_sub_scalar(lhs, rhs)
216 }
217
218 fn int_mul(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
219 B::int_mul(lhs, rhs)
220 }
221
222 fn int_mul_scalar(lhs: IntTensor<Self>, rhs: Scalar) -> IntTensor<Self> {
223 B::int_mul_scalar(lhs, rhs)
224 }
225
226 fn int_div(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
227 B::int_div(lhs, rhs)
228 }
229
230 fn int_div_scalar(lhs: IntTensor<Self>, rhs: Scalar) -> IntTensor<Self> {
231 B::int_div_scalar(lhs, rhs)
232 }
233
234 fn int_remainder(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
235 B::int_remainder(lhs, rhs)
236 }
237
238 fn int_remainder_scalar(lhs: IntTensor<Self>, rhs: Scalar) -> IntTensor<Self> {
239 B::int_remainder_scalar(lhs, rhs)
240 }
241
242 fn int_matmul(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
243 B::int_matmul(lhs, rhs)
244 }
245
246 fn int_sum(tensor: IntTensor<Self>) -> IntTensor<Self> {
247 B::int_sum(tensor)
248 }
249
250 fn int_sum_dim(tensor: IntTensor<Self>, dim: usize) -> IntTensor<Self> {
251 B::int_sum_dim(tensor, dim)
252 }
253
254 fn int_prod(tensor: IntTensor<Self>) -> IntTensor<Self> {
255 B::int_prod(tensor)
256 }
257
258 fn int_prod_dim(tensor: IntTensor<Self>, dim: usize) -> IntTensor<Self> {
259 B::int_prod_dim(tensor, dim)
260 }
261
262 fn int_mean_dim(tensor: IntTensor<Self>, dim: usize) -> IntTensor<Self> {
263 B::int_mean_dim(tensor, dim)
264 }
265
266 fn int_cumsum(tensor: IntTensor<Self>, dim: usize) -> IntTensor<Self> {
267 B::int_cumsum(tensor, dim)
268 }
269
270 fn int_cumprod(tensor: IntTensor<Self>, dim: usize) -> IntTensor<Self> {
271 B::int_cumprod(tensor, dim)
272 }
273
274 fn int_cummin(tensor: IntTensor<Self>, dim: usize) -> IntTensor<Self> {
275 B::int_cummin(tensor, dim)
276 }
277
278 fn int_cummax(tensor: IntTensor<Self>, dim: usize) -> IntTensor<Self> {
279 B::int_cummax(tensor, dim)
280 }
281
282 fn int_argmax(tensor: IntTensor<Self>, dim: usize) -> IntTensor<Self> {
283 B::int_argmax(tensor, dim)
284 }
285
286 fn int_argtopk(tensor: IntTensor<Self>, dim: usize, k: usize) -> IntTensor<Self> {
287 B::int_argtopk(tensor, dim, k)
288 }
289
290 fn int_topk(tensor: IntTensor<Self>, dim: usize, k: usize) -> IntTensor<Self> {
291 B::int_topk(tensor, dim, k)
292 }
293
294 fn int_topk_with_indices(
295 tensor: IntTensor<Self>,
296 dim: usize,
297 k: usize,
298 ) -> (IntTensor<Self>, IntTensor<Self>) {
299 B::int_topk_with_indices(tensor, dim, k)
300 }
301
302 fn int_argmin(tensor: IntTensor<Self>, dim: usize) -> IntTensor<Self> {
303 B::int_argmin(tensor, dim)
304 }
305
306 fn int_abs(tensor: IntTensor<Self>) -> IntTensor<Self> {
307 B::int_abs(tensor)
308 }
309
310 fn int_swap_dims(tensor: IntTensor<Self>, dim1: usize, dim2: usize) -> IntTensor<Self> {
311 B::int_swap_dims(tensor, dim1, dim2)
312 }
313
314 fn int_permute(tensor: IntTensor<Self>, axes: &[usize]) -> IntTensor<Self> {
315 B::int_permute(tensor, axes)
316 }
317
318 fn int_flip(tensor: IntTensor<Self>, axes: &[usize]) -> IntTensor<Self> {
319 B::int_flip(tensor, axes)
320 }
321
322 fn int_random(
323 shape: Shape,
324 distribution: burn_backend::Distribution,
325 device: &DispatchDevice,
326 dtype: IntDType,
327 ) -> IntTensor<Self> {
328 B::int_random(shape, distribution, device, dtype)
329 }
330
331 fn int_expand(tensor: IntTensor<Self>, shape: Shape) -> IntTensor<Self> {
332 B::int_expand(tensor, shape)
333 }
334
335 fn bitwise_and(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
336 B::bitwise_and(lhs, rhs)
337 }
338
339 fn bitwise_and_scalar(lhs: IntTensor<Self>, rhs: Scalar) -> IntTensor<Self> {
340 B::bitwise_and_scalar(lhs, rhs)
341 }
342
343 fn bitwise_or(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
344 B::bitwise_or(lhs, rhs)
345 }
346
347 fn bitwise_or_scalar(lhs: IntTensor<Self>, rhs: Scalar) -> IntTensor<Self> {
348 B::bitwise_or_scalar(lhs, rhs)
349 }
350
351 fn bitwise_xor(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
352 B::bitwise_xor(lhs, rhs)
353 }
354
355 fn bitwise_xor_scalar(lhs: IntTensor<Self>, rhs: Scalar) -> IntTensor<Self> {
356 B::bitwise_xor_scalar(lhs, rhs)
357 }
358
359 fn bitwise_not(tensor: IntTensor<Self>) -> IntTensor<Self> {
360 B::bitwise_not(tensor)
361 }
362
363 fn bitwise_left_shift(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
364 B::bitwise_left_shift(lhs, rhs)
365 }
366
367 fn bitwise_left_shift_scalar(lhs: IntTensor<Self>, rhs: Scalar) -> IntTensor<Self> {
368 B::bitwise_left_shift_scalar(lhs, rhs)
369 }
370
371 fn bitwise_right_shift(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
372 B::bitwise_right_shift(lhs, rhs)
373 }
374
375 fn bitwise_right_shift_scalar(lhs: IntTensor<Self>, rhs: Scalar) -> IntTensor<Self> {
376 B::bitwise_right_shift_scalar(lhs, rhs)
377 }
378
379 fn int_cast(tensor: IntTensor<Self>, dtype: IntDType) -> IntTensor<Self> {
380 B::int_cast(tensor, dtype)
381 }
382
383 fn int_unfold(
384 tensor: IntTensor<Self>,
385 dim: usize,
386 size: usize,
387 step: usize,
388 ) -> IntTensor<Self> {
389 B::int_unfold(tensor, dim, size, step)
390 }
391
392 fn int_repeat_dim(tensor: IntTensor<Self>, dim: usize, times: usize) -> IntTensor<Self> {
393 B::int_repeat_dim(tensor, dim, times)
394 }
395
396 fn int_cat(tensors: Vec<IntTensor<Self>>, dim: usize) -> IntTensor<Self> {
397 B::int_cat(tensors, dim)
398 }
399
400 fn int_not_equal(
401 lhs: IntTensor<Self>,
402 rhs: IntTensor<Self>,
403 out_dtype: BoolDType,
404 ) -> BoolTensor<Self> {
405 B::int_not_equal(lhs, rhs, out_dtype)
406 }
407
408 fn int_not_equal_elem(
409 lhs: IntTensor<Self>,
410 rhs: Scalar,
411 out_dtype: BoolDType,
412 ) -> BoolTensor<Self> {
413 B::int_not_equal_elem(lhs, rhs, out_dtype)
414 }
415
416 fn int_powi(lhs: IntTensor<Self>, rhs: IntTensor<Self>) -> IntTensor<Self> {
417 B::int_powi(lhs, rhs)
418 }
419
420 fn int_powi_scalar_impl(lhs: IntTensor<Self>, rhs: Scalar) -> IntTensor<Self> {
421 B::int_powi_scalar_impl(lhs, rhs)
422 }
423
424 fn int_clamp_min(tensor: IntTensor<Self>, min: Scalar) -> IntTensor<Self> {
425 B::int_clamp_min(tensor, min)
426 }
427
428 fn int_clamp_max(tensor: IntTensor<Self>, max: Scalar) -> IntTensor<Self> {
429 B::int_clamp_max(tensor, max)
430 }
431
432 fn int_clamp(tensor: IntTensor<Self>, min: Scalar, max: Scalar) -> IntTensor<Self> {
433 B::int_clamp(tensor, min, max)
434 }
435
436 fn int_neg(tensor: IntTensor<Self>) -> IntTensor<Self> {
437 B::int_neg(tensor)
438 }
439
440 fn int_zeros(shape: Shape, device: &DispatchDevice, dtype: IntDType) -> IntTensor<Self> {
441 B::int_zeros(shape, device, dtype)
442 }
443
444 fn int_ones(shape: Shape, device: &DispatchDevice, dtype: IntDType) -> IntTensor<Self> {
445 B::int_ones(shape, device, dtype)
446 }
447
448 fn int_full(
449 shape: Shape,
450 fill_value: Scalar,
451 device: &DispatchDevice,
452 dtype: IntDType,
453 ) -> IntTensor<Self> {
454 B::int_full(shape, fill_value, device, dtype)
455 }
456
457 fn int_mean(tensor: IntTensor<Self>) -> IntTensor<Self> {
458 B::int_mean(tensor)
459 }
460
461 fn int_max(tensor: IntTensor<Self>) -> IntTensor<Self> {
462 B::int_max(tensor)
463 }
464
465 fn int_max_dim(tensor: IntTensor<Self>, dim: usize) -> IntTensor<Self> {
466 B::int_max_dim(tensor, dim)
467 }
468
469 fn int_max_dim_with_indices(
470 tensor: IntTensor<Self>,
471 dim: usize,
472 ) -> (IntTensor<Self>, IntTensor<Self>) {
473 B::int_max_dim_with_indices(tensor, dim)
474 }
475
476 fn int_max_abs(tensor: IntTensor<Self>) -> IntTensor<Self> {
477 B::int_max_abs(tensor)
478 }
479
480 fn int_max_abs_dim(tensor: IntTensor<Self>, dim: usize) -> IntTensor<Self> {
481 B::int_max_abs_dim(tensor, dim)
482 }
483
484 fn int_min(tensor: IntTensor<Self>) -> IntTensor<Self> {
485 B::int_min(tensor)
486 }
487
488 fn int_min_dim(tensor: IntTensor<Self>, dim: usize) -> IntTensor<Self> {
489 B::int_min_dim(tensor, dim)
490 }
491
492 fn int_min_dim_with_indices(
493 tensor: IntTensor<Self>,
494 dim: usize,
495 ) -> (IntTensor<Self>, IntTensor<Self>) {
496 B::int_min_dim_with_indices(tensor, dim)
497 }
498
499 fn int_transpose(tensor: IntTensor<Self>) -> IntTensor<Self> {
500 B::int_transpose(tensor)
501 }
502
503 fn int_arange_step(
504 range: core::ops::Range<i64>,
505 step: usize,
506 device: &DispatchDevice,
507 dtype: IntDType,
508 ) -> IntTensor<Self> {
509 B::int_arange_step(range, step, device, dtype)
510 }
511
512 fn int_arange(
513 range: core::ops::Range<i64>,
514 device: &DispatchDevice,
515 dtype: IntDType,
516 ) -> IntTensor<Self> {
517 B::int_arange(range, device, dtype)
518 }
519
520 fn int_any(tensor: IntTensor<Self>, out_dtype: BoolDType) -> BoolTensor<Self> {
521 B::int_any(tensor, out_dtype)
522 }
523
524 fn int_any_dim(tensor: IntTensor<Self>, dim: usize, out_dtype: BoolDType) -> BoolTensor<Self> {
525 B::int_any_dim(tensor, dim, out_dtype)
526 }
527
528 fn int_all(tensor: IntTensor<Self>, out_dtype: BoolDType) -> BoolTensor<Self> {
529 B::int_all(tensor, out_dtype)
530 }
531
532 fn int_all_dim(tensor: IntTensor<Self>, dim: usize, out_dtype: BoolDType) -> BoolTensor<Self> {
533 B::int_all_dim(tensor, dim, out_dtype)
534 }
535
536 fn int_sign(tensor: IntTensor<Self>) -> IntTensor<Self> {
537 B::int_sign(tensor)
538 }
539
540 fn int_sort(tensor: IntTensor<Self>, dim: usize, descending: bool) -> IntTensor<Self> {
541 B::int_sort(tensor, dim, descending)
542 }
543
544 fn int_sort_with_indices(
545 tensor: IntTensor<Self>,
546 dim: usize,
547 descending: bool,
548 ) -> (IntTensor<Self>, IntTensor<Self>) {
549 B::int_sort_with_indices(tensor, dim, descending)
550 }
551
552 fn int_argsort(tensor: IntTensor<Self>, dim: usize, descending: bool) -> IntTensor<Self> {
553 B::int_argsort(tensor, dim, descending)
554 }
555}