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

burn_dispatch/ops/
bool_tensor.rs
