1use super::*;
2
3impl FloatTensorOps<Self> for Dispatch {
4 fn float_from_data(data: TensorData, device: &Device<Self>) -> FloatTensor<Self> {
5 let selection = Selection::factory(device);
6 macro_rules! call {
7 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
8 let result = <$B>::float_from_data(data, $device);
9 pack!($mode, $variant, Float, result)
10 }};
11 }
12 dispatch!(selection, call)
13 }
14
15 fn float_random( shape: Shape, distribution: Distribution, device: &Device<Self>, dtype: FloatDType, ) -> FloatTensor<Self> {
16 let selection = Selection::factory(device);
17 macro_rules! call {
18 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
19 let result = <$B>::float_random(shape, distribution, $device, dtype);
20 pack!($mode, $variant, Float, result)
21 }};
22 }
23 dispatch!(selection, call)
24 }
25
26 fn float_zeros(shape: Shape, device: &Device<Self>, dtype: FloatDType) -> FloatTensor<Self> {
27 let selection = Selection::factory(device);
28 macro_rules! call {
29 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
30 let result = <$B>::float_zeros(shape, $device, dtype);
31 pack!($mode, $variant, Float, result)
32 }};
33 }
34 dispatch!(selection, call)
35 }
36
37 fn float_ones(shape: Shape, device: &Device<Self>, dtype: FloatDType) -> FloatTensor<Self> {
38 let selection = Selection::factory(device);
39 macro_rules! call {
40 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
41 let result = <$B>::float_ones(shape, $device, dtype);
42 pack!($mode, $variant, Float, result)
43 }};
44 }
45 dispatch!(selection, call)
46 }
47
48 fn float_full( shape: Shape, fill_value: Scalar, device: &Device<Self>, dtype: FloatDType, ) -> FloatTensor<Self> {
49 let selection = Selection::factory(device);
50 macro_rules! call {
51 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
52 let result = <$B>::float_full(shape, fill_value, $device, dtype);
53 pack!($mode, $variant, Float, result)
54 }};
55 }
56 dispatch!(selection, call)
57 }
58
59 fn float_into_data( tensor: FloatTensor<Self>, ) -> impl Future<Output = Result<TensorData, ExecutionError>> + Send {
60 let selection = Selection::tensors(core::iter::once(&tensor));
61 macro_rules! call {
62 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
63 let future = <$B>::float_into_data(unpack!($mode, $variant, Float, tensor, $context));
64 let future: core::pin::Pin<Box<dyn Future<Output = Result<TensorData, ExecutionError>> + Send>> = Box::pin(async move { future.await });
65 future
66 }};
67 }
68 dispatch!(selection, call)
69 }
70
71 fn float_device(tensor: &FloatTensor<Self>) -> Device<Self> {
72 tensor.device()
73 }
74
75 fn float_to_device(tensor: FloatTensor<Self>, device: &Device<Self>) -> FloatTensor<Self> {
76 crate::transfer::float_to_device(tensor, device)
77 }
78
79 fn float_into_int(tensor: FloatTensor<Self>, out_dtype: IntDType) -> IntTensor<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>::float_into_int(unpack!($mode, $variant, Float, tensor, $context), out_dtype);
84 pack!($mode, $variant, Int, result)
85 }};
86 }
87 dispatch!(selection, call)
88 }
89
90 fn float_empty(shape: Shape, device: &Device<Self>, dtype: FloatDType) -> FloatTensor<Self> {
91 let selection = Selection::factory(device);
92 macro_rules! call {
93 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
94 let result = <$B>::float_empty(shape, $device, dtype);
95 pack!($mode, $variant, Float, result)
96 }};
97 }
98 dispatch!(selection, call)
99 }
100
101 fn float_repeat_dim(tensor: FloatTensor<Self>, dim: usize, times: usize) -> FloatTensor<Self> {
102 let selection = Selection::tensors(core::iter::once(&tensor));
103 macro_rules! call {
104 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
105 let result = <$B>::float_repeat_dim(unpack!($mode, $variant, Float, tensor, $context), dim, times);
106 pack!($mode, $variant, Float, result)
107 }};
108 }
109 dispatch!(selection, call)
110 }
111
112 fn float_add(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
113 let selection = Selection::tensors(core::iter::once(&lhs).chain(core::iter::once(&rhs)));
114 macro_rules! call {
115 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
116 let result = <$B>::float_add(unpack!($mode, $variant, Float, lhs, $context), unpack!($mode, $variant, Float, rhs, $context));
117 pack!($mode, $variant, Float, result)
118 }};
119 }
120 dispatch!(selection, call)
121 }
122
123 fn float_add_scalar(lhs: FloatTensor<Self>, rhs: Scalar) -> FloatTensor<Self> {
124 let selection = Selection::tensors(core::iter::once(&lhs));
125 macro_rules! call {
126 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
127 let result = <$B>::float_add_scalar(unpack!($mode, $variant, Float, lhs, $context), rhs);
128 pack!($mode, $variant, Float, result)
129 }};
130 }
131 dispatch!(selection, call)
132 }
133
134 fn float_clamp_min(tensor: FloatTensor<Self>, min: Scalar) -> FloatTensor<Self> {
135 let selection = Selection::tensors(core::iter::once(&tensor));
136 macro_rules! call {
137 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
138 let result = <$B>::float_clamp_min(unpack!($mode, $variant, Float, tensor, $context), min);
139 pack!($mode, $variant, Float, result)
140 }};
141 }
142 dispatch!(selection, call)
143 }
144
145 fn float_clamp_max(tensor: FloatTensor<Self>, max: Scalar) -> FloatTensor<Self> {
146 let selection = Selection::tensors(core::iter::once(&tensor));
147 macro_rules! call {
148 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
149 let result = <$B>::float_clamp_max(unpack!($mode, $variant, Float, tensor, $context), max);
150 pack!($mode, $variant, Float, result)
151 }};
152 }
153 dispatch!(selection, call)
154 }
155
156 fn float_clamp(tensor: FloatTensor<Self>, min: Scalar, max: Scalar) -> FloatTensor<Self> {
157 let selection = Selection::tensors(core::iter::once(&tensor));
158 macro_rules! call {
159 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
160 let result = <$B>::float_clamp(unpack!($mode, $variant, Float, tensor, $context), min, max);
161 pack!($mode, $variant, Float, result)
162 }};
163 }
164 dispatch!(selection, call)
165 }
166
167 fn float_sub(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
168 let selection = Selection::tensors(core::iter::once(&lhs).chain(core::iter::once(&rhs)));
169 macro_rules! call {
170 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
171 let result = <$B>::float_sub(unpack!($mode, $variant, Float, lhs, $context), unpack!($mode, $variant, Float, rhs, $context));
172 pack!($mode, $variant, Float, result)
173 }};
174 }
175 dispatch!(selection, call)
176 }
177
178 fn float_sub_scalar(lhs: FloatTensor<Self>, rhs: Scalar) -> FloatTensor<Self> {
179 let selection = Selection::tensors(core::iter::once(&lhs));
180 macro_rules! call {
181 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
182 let result = <$B>::float_sub_scalar(unpack!($mode, $variant, Float, lhs, $context), rhs);
183 pack!($mode, $variant, Float, result)
184 }};
185 }
186 dispatch!(selection, call)
187 }
188
189 fn float_mul(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
190 let selection = Selection::tensors(core::iter::once(&lhs).chain(core::iter::once(&rhs)));
191 macro_rules! call {
192 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
193 let result = <$B>::float_mul(unpack!($mode, $variant, Float, lhs, $context), unpack!($mode, $variant, Float, rhs, $context));
194 pack!($mode, $variant, Float, result)
195 }};
196 }
197 dispatch!(selection, call)
198 }
199
200 fn float_mul_scalar(lhs: FloatTensor<Self>, rhs: Scalar) -> FloatTensor<Self> {
201 let selection = Selection::tensors(core::iter::once(&lhs));
202 macro_rules! call {
203 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
204 let result = <$B>::float_mul_scalar(unpack!($mode, $variant, Float, lhs, $context), rhs);
205 pack!($mode, $variant, Float, result)
206 }};
207 }
208 dispatch!(selection, call)
209 }
210
211 fn float_div(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
212 let selection = Selection::tensors(core::iter::once(&lhs).chain(core::iter::once(&rhs)));
213 macro_rules! call {
214 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
215 let result = <$B>::float_div(unpack!($mode, $variant, Float, lhs, $context), unpack!($mode, $variant, Float, rhs, $context));
216 pack!($mode, $variant, Float, result)
217 }};
218 }
219 dispatch!(selection, call)
220 }
221
222 fn float_div_scalar(lhs: FloatTensor<Self>, rhs: Scalar) -> FloatTensor<Self> {
223 let selection = Selection::tensors(core::iter::once(&lhs));
224 macro_rules! call {
225 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
226 let result = <$B>::float_div_scalar(unpack!($mode, $variant, Float, lhs, $context), rhs);
227 pack!($mode, $variant, Float, result)
228 }};
229 }
230 dispatch!(selection, call)
231 }
232
233 fn float_remainder(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
234 let selection = Selection::tensors(core::iter::once(&lhs).chain(core::iter::once(&rhs)));
235 macro_rules! call {
236 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
237 let result = <$B>::float_remainder(unpack!($mode, $variant, Float, lhs, $context), unpack!($mode, $variant, Float, rhs, $context));
238 pack!($mode, $variant, Float, result)
239 }};
240 }
241 dispatch!(selection, call)
242 }
243
244 fn float_remainder_scalar(lhs: FloatTensor<Self>, rhs: Scalar) -> FloatTensor<Self> {
245 let selection = Selection::tensors(core::iter::once(&lhs));
246 macro_rules! call {
247 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
248 let result = <$B>::float_remainder_scalar(unpack!($mode, $variant, Float, lhs, $context), rhs);
249 pack!($mode, $variant, Float, result)
250 }};
251 }
252 dispatch!(selection, call)
253 }
254
255 fn float_matmul(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
256 let selection = Selection::tensors(core::iter::once(&lhs).chain(core::iter::once(&rhs)));
257 macro_rules! call {
258 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
259 let result = <$B>::float_matmul(unpack!($mode, $variant, Float, lhs, $context), unpack!($mode, $variant, Float, rhs, $context));
260 pack!($mode, $variant, Float, result)
261 }};
262 }
263 dispatch!(selection, call)
264 }
265
266 fn float_cross(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>, dim: usize) -> FloatTensor<Self> {
267 let selection = Selection::tensors(core::iter::once(&lhs).chain(core::iter::once(&rhs)));
268 macro_rules! call {
269 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
270 let result = <$B>::float_cross(unpack!($mode, $variant, Float, lhs, $context), unpack!($mode, $variant, Float, rhs, $context), dim);
271 pack!($mode, $variant, Float, result)
272 }};
273 }
274 dispatch!(selection, call)
275 }
276
277 fn float_neg(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
278 let selection = Selection::tensors(core::iter::once(&tensor));
279 macro_rules! call {
280 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
281 let result = <$B>::float_neg(unpack!($mode, $variant, Float, tensor, $context));
282 pack!($mode, $variant, Float, result)
283 }};
284 }
285 dispatch!(selection, call)
286 }
287
288 fn float_recip(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
289 let selection = Selection::tensors(core::iter::once(&tensor));
290 macro_rules! call {
291 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
292 let result = <$B>::float_recip(unpack!($mode, $variant, Float, tensor, $context));
293 pack!($mode, $variant, Float, result)
294 }};
295 }
296 dispatch!(selection, call)
297 }
298
299 fn float_transpose(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
300 let selection = Selection::tensors(core::iter::once(&tensor));
301 macro_rules! call {
302 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
303 let result = <$B>::float_transpose(unpack!($mode, $variant, Float, tensor, $context));
304 pack!($mode, $variant, Float, result)
305 }};
306 }
307 dispatch!(selection, call)
308 }
309
310 fn float_swap_dims(tensor: FloatTensor<Self>, dim1: usize, dim2: usize) -> FloatTensor<Self> {
311 let selection = Selection::tensors(core::iter::once(&tensor));
312 macro_rules! call {
313 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
314 let result = <$B>::float_swap_dims(unpack!($mode, $variant, Float, tensor, $context), dim1, dim2);
315 pack!($mode, $variant, Float, result)
316 }};
317 }
318 dispatch!(selection, call)
319 }
320
321 fn float_permute(tensor: FloatTensor<Self>, axes: &[usize]) -> FloatTensor<Self> {
322 let selection = Selection::tensors(core::iter::once(&tensor));
323 macro_rules! call {
324 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
325 let result = <$B>::float_permute(unpack!($mode, $variant, Float, tensor, $context), axes);
326 pack!($mode, $variant, Float, result)
327 }};
328 }
329 dispatch!(selection, call)
330 }
331
332 fn float_flip(tensor: FloatTensor<Self>, axes: &[usize]) -> FloatTensor<Self> {
333 let selection = Selection::tensors(core::iter::once(&tensor));
334 macro_rules! call {
335 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
336 let result = <$B>::float_flip(unpack!($mode, $variant, Float, tensor, $context), axes);
337 pack!($mode, $variant, Float, result)
338 }};
339 }
340 dispatch!(selection, call)
341 }
342
343 fn float_reshape(tensor: FloatTensor<Self>, shape: Shape) -> FloatTensor<Self> {
344 let selection = Selection::tensors(core::iter::once(&tensor));
345 macro_rules! call {
346 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
347 let result = <$B>::float_reshape(unpack!($mode, $variant, Float, tensor, $context), shape);
348 pack!($mode, $variant, Float, result)
349 }};
350 }
351 dispatch!(selection, call)
352 }
353
354 fn float_gather(dim: usize, tensor: FloatTensor<Self>, indices: IntTensor<Self>) -> FloatTensor<Self> {
355 let selection = Selection::tensors(core::iter::once(&tensor).chain(core::iter::once(&indices)));
356 macro_rules! call {
357 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
358 let result = <$B>::float_gather(dim, unpack!($mode, $variant, Float, tensor, $context), unpack!($mode, $variant, Int, indices, $context));
359 pack!($mode, $variant, Float, result)
360 }};
361 }
362 dispatch!(selection, call)
363 }
364
365 fn float_scatter_add( dim: usize, tensor: FloatTensor<Self>, indices: IntTensor<Self>, value: FloatTensor<Self>, ) -> FloatTensor<Self> {
366 let selection = Selection::tensors(core::iter::once(&tensor).chain(core::iter::once(&indices)).chain(core::iter::once(&value)));
367 macro_rules! call {
368 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
369 let result = <$B>::float_scatter_add(dim, unpack!($mode, $variant, Float, tensor, $context), unpack!($mode, $variant, Int, indices, $context), unpack!($mode, $variant, Float, value, $context));
370 pack!($mode, $variant, Float, result)
371 }};
372 }
373 dispatch!(selection, call)
374 }
375
376 fn float_scatter( dim: usize, tensor: FloatTensor<Self>, indices: IntTensor<Self>, value: FloatTensor<Self>, update: ruda_tensor::tensor::IndexingUpdateOp, ) -> FloatTensor<Self> {
377 let selection = Selection::tensors(core::iter::once(&tensor).chain(core::iter::once(&indices)).chain(core::iter::once(&value)));
378 macro_rules! call {
379 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
380 let result = <$B>::float_scatter(dim, unpack!($mode, $variant, Float, tensor, $context), unpack!($mode, $variant, Int, indices, $context), unpack!($mode, $variant, Float, value, $context), update);
381 pack!($mode, $variant, Float, result)
382 }};
383 }
384 dispatch!(selection, call)
385 }
386
387 fn float_scatter_nd( _data: FloatTensor<Self>, _indices: IntTensor<Self>, _values: FloatTensor<Self>, _reduction: ruda_tensor::tensor::IndexingUpdateOp, ) -> FloatTensor<Self> {
388 let selection = Selection::tensors(core::iter::once(&_data).chain(core::iter::once(&_indices)).chain(core::iter::once(&_values)));
389 macro_rules! call {
390 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
391 let result = <$B>::float_scatter_nd(unpack!($mode, $variant, Float, _data, $context), unpack!($mode, $variant, Int, _indices, $context), unpack!($mode, $variant, Float, _values, $context), _reduction);
392 pack!($mode, $variant, Float, result)
393 }};
394 }
395 dispatch!(selection, call)
396 }
397
398 fn float_gather_nd(_data: FloatTensor<Self>, _indices: IntTensor<Self>) -> FloatTensor<Self> {
399 let selection = Selection::tensors(core::iter::once(&_data).chain(core::iter::once(&_indices)));
400 macro_rules! call {
401 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
402 let result = <$B>::float_gather_nd(unpack!($mode, $variant, Float, _data, $context), unpack!($mode, $variant, Int, _indices, $context));
403 pack!($mode, $variant, Float, result)
404 }};
405 }
406 dispatch!(selection, call)
407 }
408
409 fn float_select(tensor: FloatTensor<Self>, dim: usize, indices: IntTensor<Self>) -> FloatTensor<Self> {
410 let selection = Selection::tensors(core::iter::once(&tensor).chain(core::iter::once(&indices)));
411 macro_rules! call {
412 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
413 let result = <$B>::float_select(unpack!($mode, $variant, Float, tensor, $context), dim, unpack!($mode, $variant, Int, indices, $context));
414 pack!($mode, $variant, Float, result)
415 }};
416 }
417 dispatch!(selection, call)
418 }
419
420 fn float_select_add( tensor: FloatTensor<Self>, dim: usize, indices: IntTensor<Self>, value: FloatTensor<Self>, ) -> FloatTensor<Self> {
421 let selection = Selection::tensors(core::iter::once(&tensor).chain(core::iter::once(&indices)).chain(core::iter::once(&value)));
422 macro_rules! call {
423 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
424 let result = <$B>::float_select_add(unpack!($mode, $variant, Float, tensor, $context), dim, unpack!($mode, $variant, Int, indices, $context), unpack!($mode, $variant, Float, value, $context));
425 pack!($mode, $variant, Float, result)
426 }};
427 }
428 dispatch!(selection, call)
429 }
430
431 fn float_select_assign( tensor: FloatTensor<Self>, dim: usize, indices: IntTensor<Self>, value: FloatTensor<Self>, update: ruda_tensor::tensor::IndexingUpdateOp, ) -> FloatTensor<Self> {
432 let selection = Selection::tensors(core::iter::once(&tensor).chain(core::iter::once(&indices)).chain(core::iter::once(&value)));
433 macro_rules! call {
434 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
435 let result = <$B>::float_select_assign(unpack!($mode, $variant, Float, tensor, $context), dim, unpack!($mode, $variant, Int, indices, $context), unpack!($mode, $variant, Float, value, $context), update);
436 pack!($mode, $variant, Float, result)
437 }};
438 }
439 dispatch!(selection, call)
440 }
441
442 fn float_slice(tensor: FloatTensor<Self>, slices: &[Slice]) -> FloatTensor<Self> {
443 let selection = Selection::tensors(core::iter::once(&tensor));
444 macro_rules! call {
445 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
446 let result = <$B>::float_slice(unpack!($mode, $variant, Float, tensor, $context), slices);
447 pack!($mode, $variant, Float, result)
448 }};
449 }
450 dispatch!(selection, call)
451 }
452
453 fn float_slice_assign( tensor: FloatTensor<Self>, slices: &[Slice], value: FloatTensor<Self>, ) -> FloatTensor<Self> {
454 let selection = Selection::tensors(core::iter::once(&tensor).chain(core::iter::once(&value)));
455 macro_rules! call {
456 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
457 let result = <$B>::float_slice_assign(unpack!($mode, $variant, Float, tensor, $context), slices, unpack!($mode, $variant, Float, value, $context));
458 pack!($mode, $variant, Float, result)
459 }};
460 }
461 dispatch!(selection, call)
462 }
463
464 fn float_mask_where( tensor: FloatTensor<Self>, mask: BoolTensor<Self>, value: FloatTensor<Self>, ) -> FloatTensor<Self> {
465 let selection = Selection::tensors(core::iter::once(&tensor).chain(core::iter::once(&mask)).chain(core::iter::once(&value)));
466 macro_rules! call {
467 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
468 let result = <$B>::float_mask_where(unpack!($mode, $variant, Float, tensor, $context), unpack!($mode, $variant, Bool, mask, $context), unpack!($mode, $variant, Float, value, $context));
469 pack!($mode, $variant, Float, result)
470 }};
471 }
472 dispatch!(selection, call)
473 }
474
475 fn float_mask_fill( tensor: FloatTensor<Self>, mask: BoolTensor<Self>, value: Scalar, ) -> FloatTensor<Self> {
476 let selection = Selection::tensors(core::iter::once(&tensor).chain(core::iter::once(&mask)));
477 macro_rules! call {
478 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
479 let result = <$B>::float_mask_fill(unpack!($mode, $variant, Float, tensor, $context), unpack!($mode, $variant, Bool, mask, $context), value);
480 pack!($mode, $variant, Float, result)
481 }};
482 }
483 dispatch!(selection, call)
484 }
485
486 fn float_equal(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>, out_dtype: BoolDType) -> BoolTensor<Self> {
487 let selection = Selection::tensors(core::iter::once(&lhs).chain(core::iter::once(&rhs)));
488 macro_rules! call {
489 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
490 let result = <$B>::float_equal(unpack!($mode, $variant, Float, lhs, $context), unpack!($mode, $variant, Float, rhs, $context), out_dtype);
491 pack!($mode, $variant, Bool, result)
492 }};
493 }
494 dispatch!(selection, call)
495 }
496
497 fn float_not_equal( lhs: FloatTensor<Self>, rhs: FloatTensor<Self>, out_dtype: BoolDType, ) -> BoolTensor<Self> {
498 let selection = Selection::tensors(core::iter::once(&lhs).chain(core::iter::once(&rhs)));
499 macro_rules! call {
500 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
501 let result = <$B>::float_not_equal(unpack!($mode, $variant, Float, lhs, $context), unpack!($mode, $variant, Float, rhs, $context), out_dtype);
502 pack!($mode, $variant, Bool, result)
503 }};
504 }
505 dispatch!(selection, call)
506 }
507
508 fn float_equal_elem(lhs: FloatTensor<Self>, rhs: Scalar, out_dtype: BoolDType) -> BoolTensor<Self> {
509 let selection = Selection::tensors(core::iter::once(&lhs));
510 macro_rules! call {
511 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
512 let result = <$B>::float_equal_elem(unpack!($mode, $variant, Float, lhs, $context), rhs, out_dtype);
513 pack!($mode, $variant, Bool, result)
514 }};
515 }
516 dispatch!(selection, call)
517 }
518
519 fn float_not_equal_elem( lhs: FloatTensor<Self>, rhs: Scalar, out_dtype: BoolDType, ) -> BoolTensor<Self> {
520 let selection = Selection::tensors(core::iter::once(&lhs));
521 macro_rules! call {
522 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
523 let result = <$B>::float_not_equal_elem(unpack!($mode, $variant, Float, lhs, $context), rhs, out_dtype);
524 pack!($mode, $variant, Bool, result)
525 }};
526 }
527 dispatch!(selection, call)
528 }
529
530 fn float_greater( lhs: FloatTensor<Self>, rhs: FloatTensor<Self>, out_dtype: BoolDType, ) -> BoolTensor<Self> {
531 let selection = Selection::tensors(core::iter::once(&lhs).chain(core::iter::once(&rhs)));
532 macro_rules! call {
533 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
534 let result = <$B>::float_greater(unpack!($mode, $variant, Float, lhs, $context), unpack!($mode, $variant, Float, rhs, $context), out_dtype);
535 pack!($mode, $variant, Bool, result)
536 }};
537 }
538 dispatch!(selection, call)
539 }
540
541 fn float_greater_elem(lhs: FloatTensor<Self>, rhs: Scalar, out_dtype: BoolDType) -> BoolTensor<Self> {
542 let selection = Selection::tensors(core::iter::once(&lhs));
543 macro_rules! call {
544 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
545 let result = <$B>::float_greater_elem(unpack!($mode, $variant, Float, lhs, $context), rhs, out_dtype);
546 pack!($mode, $variant, Bool, result)
547 }};
548 }
549 dispatch!(selection, call)
550 }
551
552 fn float_greater_equal( lhs: FloatTensor<Self>, rhs: FloatTensor<Self>, out_dtype: BoolDType, ) -> BoolTensor<Self> {
553 let selection = Selection::tensors(core::iter::once(&lhs).chain(core::iter::once(&rhs)));
554 macro_rules! call {
555 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
556 let result = <$B>::float_greater_equal(unpack!($mode, $variant, Float, lhs, $context), unpack!($mode, $variant, Float, rhs, $context), out_dtype);
557 pack!($mode, $variant, Bool, result)
558 }};
559 }
560 dispatch!(selection, call)
561 }
562
563 fn float_greater_equal_elem( lhs: FloatTensor<Self>, rhs: Scalar, out_dtype: BoolDType, ) -> BoolTensor<Self> {
564 let selection = Selection::tensors(core::iter::once(&lhs));
565 macro_rules! call {
566 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
567 let result = <$B>::float_greater_equal_elem(unpack!($mode, $variant, Float, lhs, $context), rhs, out_dtype);
568 pack!($mode, $variant, Bool, result)
569 }};
570 }
571 dispatch!(selection, call)
572 }
573
574 fn float_lower(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>, out_dtype: BoolDType) -> BoolTensor<Self> {
575 let selection = Selection::tensors(core::iter::once(&lhs).chain(core::iter::once(&rhs)));
576 macro_rules! call {
577 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
578 let result = <$B>::float_lower(unpack!($mode, $variant, Float, lhs, $context), unpack!($mode, $variant, Float, rhs, $context), out_dtype);
579 pack!($mode, $variant, Bool, result)
580 }};
581 }
582 dispatch!(selection, call)
583 }
584
585 fn float_lower_elem(lhs: FloatTensor<Self>, rhs: Scalar, out_dtype: BoolDType) -> BoolTensor<Self> {
586 let selection = Selection::tensors(core::iter::once(&lhs));
587 macro_rules! call {
588 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
589 let result = <$B>::float_lower_elem(unpack!($mode, $variant, Float, lhs, $context), rhs, out_dtype);
590 pack!($mode, $variant, Bool, result)
591 }};
592 }
593 dispatch!(selection, call)
594 }
595
596 fn float_lower_equal( lhs: FloatTensor<Self>, rhs: FloatTensor<Self>, out_dtype: BoolDType, ) -> BoolTensor<Self> {
597 let selection = Selection::tensors(core::iter::once(&lhs).chain(core::iter::once(&rhs)));
598 macro_rules! call {
599 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
600 let result = <$B>::float_lower_equal(unpack!($mode, $variant, Float, lhs, $context), unpack!($mode, $variant, Float, rhs, $context), out_dtype);
601 pack!($mode, $variant, Bool, result)
602 }};
603 }
604 dispatch!(selection, call)
605 }
606
607 fn float_lower_equal_elem( lhs: FloatTensor<Self>, rhs: Scalar, out_dtype: BoolDType, ) -> BoolTensor<Self> {
608 let selection = Selection::tensors(core::iter::once(&lhs));
609 macro_rules! call {
610 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
611 let result = <$B>::float_lower_equal_elem(unpack!($mode, $variant, Float, lhs, $context), rhs, out_dtype);
612 pack!($mode, $variant, Bool, result)
613 }};
614 }
615 dispatch!(selection, call)
616 }
617
618 fn float_detach(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
619 let selection = Selection::tensors(core::iter::once(&tensor));
620 macro_rules! call {
621 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
622 let result = <$B>::float_detach(unpack!($mode, $variant, Float, tensor, $context));
623 pack!($mode, $variant, Float, result)
624 }};
625 }
626 dispatch!(selection, call)
627 }
628
629 fn float_set_require_grad(tensor: FloatTensor<Self>, _require_grad: bool) -> FloatTensor<Self> {
630 let selection = Selection::tensors(core::iter::once(&tensor));
631 macro_rules! call {
632 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
633 let result = <$B>::float_set_require_grad(unpack!($mode, $variant, Float, tensor, $context), _require_grad);
634 pack!($mode, $variant, Float, result)
635 }};
636 }
637 dispatch!(selection, call)
638 }
639
640 fn float_is_require_grad(_tensor: &FloatTensor<Self>) -> bool {
641 let selection = Selection::tensors(core::iter::once(_tensor));
642 macro_rules! call {
643 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
644 let result = <$B>::float_is_require_grad(borrow_tensor!($mode, $variant, Float, _tensor));
645 result
646 }};
647 }
648 dispatch!(selection, call)
649 }
650
651 fn float_is_tracked(tensor: &FloatTensor<Self>) -> bool {
652 let selection = Selection::tensors(core::iter::once(tensor));
653 macro_rules! call {
654 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
655 let result = <$B>::float_is_tracked(borrow_tensor!($mode, $variant, Float, tensor));
656 result
657 }};
658 }
659 dispatch!(selection, call)
660 }
661
662 fn float_gradient_checkpointing_strategy(_tensor: &FloatTensor<Self>) -> Option<ruda_tensor::GradientCheckpointingStrategy> {
663 match _tensor.context() { AutodiffContext::Disabled => None, AutodiffContext::Enabled(strategy) => Some(strategy) }
664 }
665
666 fn float_autodiff_context(tensor: &FloatTensor<Self>) -> ruda_tensor::AutodiffContext {
667 tensor.context()
668 }
669
670 fn float_with_autodiff_context(tensor: FloatTensor<Self>, context: ruda_tensor::AutodiffContext) -> FloatTensor<Self> {
671 tensor.with_context(context)
672 }
673
674 fn float_with_gradient_checkpointing_strategy( _tensor: FloatTensor<Self>, _strategy: ruda_tensor::GradientCheckpointingStrategy, ) -> FloatTensor<Self> {
675 assert!(matches!(_tensor.context(), AutodiffContext::Enabled(_)), "Checkpointing requires autodiff");
676 _tensor.with_context(AutodiffContext::Enabled(_strategy))
677 }
678
679 fn float_sum(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
680 let selection = Selection::tensors(core::iter::once(&tensor));
681 macro_rules! call {
682 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
683 let result = <$B>::float_sum(unpack!($mode, $variant, Float, tensor, $context));
684 pack!($mode, $variant, Float, result)
685 }};
686 }
687 dispatch!(selection, call)
688 }
689
690 fn float_sum_dim(tensor: FloatTensor<Self>, dim: usize) -> FloatTensor<Self> {
691 let selection = Selection::tensors(core::iter::once(&tensor));
692 macro_rules! call {
693 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
694 let result = <$B>::float_sum_dim(unpack!($mode, $variant, Float, tensor, $context), dim);
695 pack!($mode, $variant, Float, result)
696 }};
697 }
698 dispatch!(selection, call)
699 }
700
701 fn float_prod(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
702 let selection = Selection::tensors(core::iter::once(&tensor));
703 macro_rules! call {
704 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
705 let result = <$B>::float_prod(unpack!($mode, $variant, Float, tensor, $context));
706 pack!($mode, $variant, Float, result)
707 }};
708 }
709 dispatch!(selection, call)
710 }
711
712 fn float_prod_dim(tensor: FloatTensor<Self>, dim: usize) -> FloatTensor<Self> {
713 let selection = Selection::tensors(core::iter::once(&tensor));
714 macro_rules! call {
715 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
716 let result = <$B>::float_prod_dim(unpack!($mode, $variant, Float, tensor, $context), dim);
717 pack!($mode, $variant, Float, result)
718 }};
719 }
720 dispatch!(selection, call)
721 }
722
723 fn float_mean(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
724 let selection = Selection::tensors(core::iter::once(&tensor));
725 macro_rules! call {
726 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
727 let result = <$B>::float_mean(unpack!($mode, $variant, Float, tensor, $context));
728 pack!($mode, $variant, Float, result)
729 }};
730 }
731 dispatch!(selection, call)
732 }
733
734 fn float_mean_dim(tensor: FloatTensor<Self>, dim: usize) -> FloatTensor<Self> {
735 let selection = Selection::tensors(core::iter::once(&tensor));
736 macro_rules! call {
737 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
738 let result = <$B>::float_mean_dim(unpack!($mode, $variant, Float, tensor, $context), dim);
739 pack!($mode, $variant, Float, result)
740 }};
741 }
742 dispatch!(selection, call)
743 }
744
745 fn float_cumsum(tensor: FloatTensor<Self>, dim: usize) -> FloatTensor<Self> {
746 let selection = Selection::tensors(core::iter::once(&tensor));
747 macro_rules! call {
748 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
749 let result = <$B>::float_cumsum(unpack!($mode, $variant, Float, tensor, $context), dim);
750 pack!($mode, $variant, Float, result)
751 }};
752 }
753 dispatch!(selection, call)
754 }
755
756 fn float_cumprod(tensor: FloatTensor<Self>, dim: usize) -> FloatTensor<Self> {
757 let selection = Selection::tensors(core::iter::once(&tensor));
758 macro_rules! call {
759 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
760 let result = <$B>::float_cumprod(unpack!($mode, $variant, Float, tensor, $context), dim);
761 pack!($mode, $variant, Float, result)
762 }};
763 }
764 dispatch!(selection, call)
765 }
766
767 fn float_cummin(tensor: FloatTensor<Self>, dim: usize) -> FloatTensor<Self> {
768 let selection = Selection::tensors(core::iter::once(&tensor));
769 macro_rules! call {
770 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
771 let result = <$B>::float_cummin(unpack!($mode, $variant, Float, tensor, $context), dim);
772 pack!($mode, $variant, Float, result)
773 }};
774 }
775 dispatch!(selection, call)
776 }
777
778 fn float_cummax(tensor: FloatTensor<Self>, dim: usize) -> FloatTensor<Self> {
779 let selection = Selection::tensors(core::iter::once(&tensor));
780 macro_rules! call {
781 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
782 let result = <$B>::float_cummax(unpack!($mode, $variant, Float, tensor, $context), dim);
783 pack!($mode, $variant, Float, result)
784 }};
785 }
786 dispatch!(selection, call)
787 }
788
789 fn float_cast(tensor: FloatTensor<Self>, dtype: FloatDType) -> FloatTensor<Self> {
790 let selection = Selection::tensors(core::iter::once(&tensor));
791 macro_rules! call {
792 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
793 let result = <$B>::float_cast(unpack!($mode, $variant, Float, tensor, $context), dtype);
794 pack!($mode, $variant, Float, result)
795 }};
796 }
797 dispatch!(selection, call)
798 }
799
800 fn float_exp(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
801 let selection = Selection::tensors(core::iter::once(&tensor));
802 macro_rules! call {
803 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
804 let result = <$B>::float_exp(unpack!($mode, $variant, Float, tensor, $context));
805 pack!($mode, $variant, Float, result)
806 }};
807 }
808 dispatch!(selection, call)
809 }
810
811 fn float_log(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
812 let selection = Selection::tensors(core::iter::once(&tensor));
813 macro_rules! call {
814 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
815 let result = <$B>::float_log(unpack!($mode, $variant, Float, tensor, $context));
816 pack!($mode, $variant, Float, result)
817 }};
818 }
819 dispatch!(selection, call)
820 }
821
822 fn float_log1p(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
823 let selection = Selection::tensors(core::iter::once(&tensor));
824 macro_rules! call {
825 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
826 let result = <$B>::float_log1p(unpack!($mode, $variant, Float, tensor, $context));
827 pack!($mode, $variant, Float, result)
828 }};
829 }
830 dispatch!(selection, call)
831 }
832
833 fn float_powf(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
834 let selection = Selection::tensors(core::iter::once(&lhs).chain(core::iter::once(&rhs)));
835 macro_rules! call {
836 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
837 let result = <$B>::float_powf(unpack!($mode, $variant, Float, lhs, $context), unpack!($mode, $variant, Float, rhs, $context));
838 pack!($mode, $variant, Float, result)
839 }};
840 }
841 dispatch!(selection, call)
842 }
843
844 fn float_powi(lhs: FloatTensor<Self>, rhs: IntTensor<Self>) -> FloatTensor<Self> {
845 let selection = Selection::tensors(core::iter::once(&lhs).chain(core::iter::once(&rhs)));
846 macro_rules! call {
847 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
848 let result = <$B>::float_powi(unpack!($mode, $variant, Float, lhs, $context), unpack!($mode, $variant, Int, rhs, $context));
849 pack!($mode, $variant, Float, result)
850 }};
851 }
852 dispatch!(selection, call)
853 }
854
855 fn float_powi_scalar(lhs: FloatTensor<Self>, rhs: Scalar) -> FloatTensor<Self> {
856 let selection = Selection::tensors(core::iter::once(&lhs));
857 macro_rules! call {
858 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
859 let result = <$B>::float_powi_scalar(unpack!($mode, $variant, Float, lhs, $context), rhs);
860 pack!($mode, $variant, Float, result)
861 }};
862 }
863 dispatch!(selection, call)
864 }
865
866 fn float_powi_scalar_impl(lhs: FloatTensor<Self>, rhs: Scalar) -> FloatTensor<Self> {
867 let selection = Selection::tensors(core::iter::once(&lhs));
868 macro_rules! call {
869 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
870 let result = <$B>::float_powi_scalar_impl(unpack!($mode, $variant, Float, lhs, $context), rhs);
871 pack!($mode, $variant, Float, result)
872 }};
873 }
874 dispatch!(selection, call)
875 }
876
877 fn float_powf_scalar(tensor: FloatTensor<Self>, value: Scalar) -> FloatTensor<Self> {
878 let selection = Selection::tensors(core::iter::once(&tensor));
879 macro_rules! call {
880 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
881 let result = <$B>::float_powf_scalar(unpack!($mode, $variant, Float, tensor, $context), value);
882 pack!($mode, $variant, Float, result)
883 }};
884 }
885 dispatch!(selection, call)
886 }
887
888 fn float_powf_scalar_impl(tensor: FloatTensor<Self>, value: Scalar) -> FloatTensor<Self> {
889 let selection = Selection::tensors(core::iter::once(&tensor));
890 macro_rules! call {
891 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
892 let result = <$B>::float_powf_scalar_impl(unpack!($mode, $variant, Float, tensor, $context), value);
893 pack!($mode, $variant, Float, result)
894 }};
895 }
896 dispatch!(selection, call)
897 }
898
899 fn float_sqrt(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
900 let selection = Selection::tensors(core::iter::once(&tensor));
901 macro_rules! call {
902 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
903 let result = <$B>::float_sqrt(unpack!($mode, $variant, Float, tensor, $context));
904 pack!($mode, $variant, Float, result)
905 }};
906 }
907 dispatch!(selection, call)
908 }
909
910 fn float_hypot(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
911 let selection = Selection::tensors(core::iter::once(&lhs).chain(core::iter::once(&rhs)));
912 macro_rules! call {
913 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
914 let result = <$B>::float_hypot(unpack!($mode, $variant, Float, lhs, $context), unpack!($mode, $variant, Float, rhs, $context));
915 pack!($mode, $variant, Float, result)
916 }};
917 }
918 dispatch!(selection, call)
919 }
920
921 fn float_rsqrt(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
922 let selection = Selection::tensors(core::iter::once(&tensor));
923 macro_rules! call {
924 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
925 let result = <$B>::float_rsqrt(unpack!($mode, $variant, Float, tensor, $context));
926 pack!($mode, $variant, Float, result)
927 }};
928 }
929 dispatch!(selection, call)
930 }
931
932 fn float_abs(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
933 let selection = Selection::tensors(core::iter::once(&tensor));
934 macro_rules! call {
935 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
936 let result = <$B>::float_abs(unpack!($mode, $variant, Float, tensor, $context));
937 pack!($mode, $variant, Float, result)
938 }};
939 }
940 dispatch!(selection, call)
941 }
942
943 fn float_cos(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
944 let selection = Selection::tensors(core::iter::once(&tensor));
945 macro_rules! call {
946 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
947 let result = <$B>::float_cos(unpack!($mode, $variant, Float, tensor, $context));
948 pack!($mode, $variant, Float, result)
949 }};
950 }
951 dispatch!(selection, call)
952 }
953
954 fn float_sin(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
955 let selection = Selection::tensors(core::iter::once(&tensor));
956 macro_rules! call {
957 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
958 let result = <$B>::float_sin(unpack!($mode, $variant, Float, tensor, $context));
959 pack!($mode, $variant, Float, result)
960 }};
961 }
962 dispatch!(selection, call)
963 }
964
965 fn float_tan(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
966 let selection = Selection::tensors(core::iter::once(&tensor));
967 macro_rules! call {
968 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
969 let result = <$B>::float_tan(unpack!($mode, $variant, Float, tensor, $context));
970 pack!($mode, $variant, Float, result)
971 }};
972 }
973 dispatch!(selection, call)
974 }
975
976 fn float_cosh(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
977 let selection = Selection::tensors(core::iter::once(&tensor));
978 macro_rules! call {
979 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
980 let result = <$B>::float_cosh(unpack!($mode, $variant, Float, tensor, $context));
981 pack!($mode, $variant, Float, result)
982 }};
983 }
984 dispatch!(selection, call)
985 }
986
987 fn float_sinh(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
988 let selection = Selection::tensors(core::iter::once(&tensor));
989 macro_rules! call {
990 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
991 let result = <$B>::float_sinh(unpack!($mode, $variant, Float, tensor, $context));
992 pack!($mode, $variant, Float, result)
993 }};
994 }
995 dispatch!(selection, call)
996 }
997
998 fn float_tanh(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
999 let selection = Selection::tensors(core::iter::once(&tensor));
1000 macro_rules! call {
1001 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1002 let result = <$B>::float_tanh(unpack!($mode, $variant, Float, tensor, $context));
1003 pack!($mode, $variant, Float, result)
1004 }};
1005 }
1006 dispatch!(selection, call)
1007 }
1008
1009 fn float_acos(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
1010 let selection = Selection::tensors(core::iter::once(&tensor));
1011 macro_rules! call {
1012 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1013 let result = <$B>::float_acos(unpack!($mode, $variant, Float, tensor, $context));
1014 pack!($mode, $variant, Float, result)
1015 }};
1016 }
1017 dispatch!(selection, call)
1018 }
1019
1020 fn float_acosh(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
1021 let selection = Selection::tensors(core::iter::once(&tensor));
1022 macro_rules! call {
1023 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1024 let result = <$B>::float_acosh(unpack!($mode, $variant, Float, tensor, $context));
1025 pack!($mode, $variant, Float, result)
1026 }};
1027 }
1028 dispatch!(selection, call)
1029 }
1030
1031 fn float_asin(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
1032 let selection = Selection::tensors(core::iter::once(&tensor));
1033 macro_rules! call {
1034 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1035 let result = <$B>::float_asin(unpack!($mode, $variant, Float, tensor, $context));
1036 pack!($mode, $variant, Float, result)
1037 }};
1038 }
1039 dispatch!(selection, call)
1040 }
1041
1042 fn float_asinh(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
1043 let selection = Selection::tensors(core::iter::once(&tensor));
1044 macro_rules! call {
1045 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1046 let result = <$B>::float_asinh(unpack!($mode, $variant, Float, tensor, $context));
1047 pack!($mode, $variant, Float, result)
1048 }};
1049 }
1050 dispatch!(selection, call)
1051 }
1052
1053 fn float_atan(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
1054 let selection = Selection::tensors(core::iter::once(&tensor));
1055 macro_rules! call {
1056 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1057 let result = <$B>::float_atan(unpack!($mode, $variant, Float, tensor, $context));
1058 pack!($mode, $variant, Float, result)
1059 }};
1060 }
1061 dispatch!(selection, call)
1062 }
1063
1064 fn float_atanh(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
1065 let selection = Selection::tensors(core::iter::once(&tensor));
1066 macro_rules! call {
1067 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1068 let result = <$B>::float_atanh(unpack!($mode, $variant, Float, tensor, $context));
1069 pack!($mode, $variant, Float, result)
1070 }};
1071 }
1072 dispatch!(selection, call)
1073 }
1074
1075 fn float_atan2(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
1076 let selection = Selection::tensors(core::iter::once(&lhs).chain(core::iter::once(&rhs)));
1077 macro_rules! call {
1078 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1079 let result = <$B>::float_atan2(unpack!($mode, $variant, Float, lhs, $context), unpack!($mode, $variant, Float, rhs, $context));
1080 pack!($mode, $variant, Float, result)
1081 }};
1082 }
1083 dispatch!(selection, call)
1084 }
1085
1086 fn float_round(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
1087 let selection = Selection::tensors(core::iter::once(&tensor));
1088 macro_rules! call {
1089 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1090 let result = <$B>::float_round(unpack!($mode, $variant, Float, tensor, $context));
1091 pack!($mode, $variant, Float, result)
1092 }};
1093 }
1094 dispatch!(selection, call)
1095 }
1096
1097 fn float_floor(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
1098 let selection = Selection::tensors(core::iter::once(&tensor));
1099 macro_rules! call {
1100 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1101 let result = <$B>::float_floor(unpack!($mode, $variant, Float, tensor, $context));
1102 pack!($mode, $variant, Float, result)
1103 }};
1104 }
1105 dispatch!(selection, call)
1106 }
1107
1108 fn float_ceil(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
1109 let selection = Selection::tensors(core::iter::once(&tensor));
1110 macro_rules! call {
1111 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1112 let result = <$B>::float_ceil(unpack!($mode, $variant, Float, tensor, $context));
1113 pack!($mode, $variant, Float, result)
1114 }};
1115 }
1116 dispatch!(selection, call)
1117 }
1118
1119 fn float_trunc(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
1120 let selection = Selection::tensors(core::iter::once(&tensor));
1121 macro_rules! call {
1122 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1123 let result = <$B>::float_trunc(unpack!($mode, $variant, Float, tensor, $context));
1124 pack!($mode, $variant, Float, result)
1125 }};
1126 }
1127 dispatch!(selection, call)
1128 }
1129
1130 fn float_erf(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
1131 let selection = Selection::tensors(core::iter::once(&tensor));
1132 macro_rules! call {
1133 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1134 let result = <$B>::float_erf(unpack!($mode, $variant, Float, tensor, $context));
1135 pack!($mode, $variant, Float, result)
1136 }};
1137 }
1138 dispatch!(selection, call)
1139 }
1140
1141 fn float_cat(tensors: Vec<FloatTensor<Self>>, dim: usize) -> FloatTensor<Self> {
1142 let selection = Selection::tensors(tensors.iter());
1143 macro_rules! call {
1144 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1145 let result = <$B>::float_cat(tensors.into_iter().map(|value| unpack!($mode, $variant, Float, value, $context)).collect(), dim);
1146 pack!($mode, $variant, Float, result)
1147 }};
1148 }
1149 dispatch!(selection, call)
1150 }
1151
1152 fn float_argmax(tensor: FloatTensor<Self>, dim: usize, out_dtype: IntDType) -> IntTensor<Self> {
1153 let selection = Selection::tensors(core::iter::once(&tensor));
1154 macro_rules! call {
1155 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1156 let result = <$B>::float_argmax(unpack!($mode, $variant, Float, tensor, $context), dim, out_dtype);
1157 pack!($mode, $variant, Int, result)
1158 }};
1159 }
1160 dispatch!(selection, call)
1161 }
1162
1163 fn float_argtopk( tensor: FloatTensor<Self>, dim: usize, k: usize, out_dtype: IntDType, ) -> IntTensor<Self> {
1164 let selection = Selection::tensors(core::iter::once(&tensor));
1165 macro_rules! call {
1166 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1167 let result = <$B>::float_argtopk(unpack!($mode, $variant, Float, tensor, $context), dim, k, out_dtype);
1168 pack!($mode, $variant, Int, result)
1169 }};
1170 }
1171 dispatch!(selection, call)
1172 }
1173
1174 fn float_topk(tensor: FloatTensor<Self>, dim: usize, k: usize) -> FloatTensor<Self> {
1175 let selection = Selection::tensors(core::iter::once(&tensor));
1176 macro_rules! call {
1177 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1178 let result = <$B>::float_topk(unpack!($mode, $variant, Float, tensor, $context), dim, k);
1179 pack!($mode, $variant, Float, result)
1180 }};
1181 }
1182 dispatch!(selection, call)
1183 }
1184
1185 fn float_argmin(tensor: FloatTensor<Self>, dim: usize, out_dtype: IntDType) -> IntTensor<Self> {
1186 let selection = Selection::tensors(core::iter::once(&tensor));
1187 macro_rules! call {
1188 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1189 let result = <$B>::float_argmin(unpack!($mode, $variant, Float, tensor, $context), dim, out_dtype);
1190 pack!($mode, $variant, Int, result)
1191 }};
1192 }
1193 dispatch!(selection, call)
1194 }
1195
1196 fn float_max(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
1197 let selection = Selection::tensors(core::iter::once(&tensor));
1198 macro_rules! call {
1199 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1200 let result = <$B>::float_max(unpack!($mode, $variant, Float, tensor, $context));
1201 pack!($mode, $variant, Float, result)
1202 }};
1203 }
1204 dispatch!(selection, call)
1205 }
1206
1207 fn float_max_dim(tensor: FloatTensor<Self>, dim: usize) -> FloatTensor<Self> {
1208 let selection = Selection::tensors(core::iter::once(&tensor));
1209 macro_rules! call {
1210 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1211 let result = <$B>::float_max_dim(unpack!($mode, $variant, Float, tensor, $context), dim);
1212 pack!($mode, $variant, Float, result)
1213 }};
1214 }
1215 dispatch!(selection, call)
1216 }
1217
1218 fn float_max_dim_with_indices( tensor: FloatTensor<Self>, dim: usize, indices_dtype: IntDType, ) -> (FloatTensor<Self>, IntTensor<Self>) {
1219 let selection = Selection::tensors(core::iter::once(&tensor));
1220 macro_rules! call {
1221 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1222 let result = <$B>::float_max_dim_with_indices(unpack!($mode, $variant, Float, tensor, $context), dim, indices_dtype);
1223 { let (value0, value1) = result; (pack!($mode, $variant, Float, value0), pack!($mode, $variant, Int, value1)) }
1224 }};
1225 }
1226 dispatch!(selection, call)
1227 }
1228
1229 fn float_min(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
1230 let selection = Selection::tensors(core::iter::once(&tensor));
1231 macro_rules! call {
1232 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1233 let result = <$B>::float_min(unpack!($mode, $variant, Float, tensor, $context));
1234 pack!($mode, $variant, Float, result)
1235 }};
1236 }
1237 dispatch!(selection, call)
1238 }
1239
1240 fn float_min_dim(tensor: FloatTensor<Self>, dim: usize) -> FloatTensor<Self> {
1241 let selection = Selection::tensors(core::iter::once(&tensor));
1242 macro_rules! call {
1243 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1244 let result = <$B>::float_min_dim(unpack!($mode, $variant, Float, tensor, $context), dim);
1245 pack!($mode, $variant, Float, result)
1246 }};
1247 }
1248 dispatch!(selection, call)
1249 }
1250
1251 fn float_min_dim_with_indices( tensor: FloatTensor<Self>, dim: usize, indices_dtype: IntDType, ) -> (FloatTensor<Self>, IntTensor<Self>) {
1252 let selection = Selection::tensors(core::iter::once(&tensor));
1253 macro_rules! call {
1254 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1255 let result = <$B>::float_min_dim_with_indices(unpack!($mode, $variant, Float, tensor, $context), dim, indices_dtype);
1256 { let (value0, value1) = result; (pack!($mode, $variant, Float, value0), pack!($mode, $variant, Int, value1)) }
1257 }};
1258 }
1259 dispatch!(selection, call)
1260 }
1261
1262 fn float_max_abs(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
1263 let selection = Selection::tensors(core::iter::once(&tensor));
1264 macro_rules! call {
1265 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1266 let result = <$B>::float_max_abs(unpack!($mode, $variant, Float, tensor, $context));
1267 pack!($mode, $variant, Float, result)
1268 }};
1269 }
1270 dispatch!(selection, call)
1271 }
1272
1273 fn float_max_abs_dim(tensor: FloatTensor<Self>, dim: usize) -> FloatTensor<Self> {
1274 let selection = Selection::tensors(core::iter::once(&tensor));
1275 macro_rules! call {
1276 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1277 let result = <$B>::float_max_abs_dim(unpack!($mode, $variant, Float, tensor, $context), dim);
1278 pack!($mode, $variant, Float, result)
1279 }};
1280 }
1281 dispatch!(selection, call)
1282 }
1283
1284 fn float_any(tensor: FloatTensor<Self>, out_dtype: BoolDType) -> BoolTensor<Self> {
1285 let selection = Selection::tensors(core::iter::once(&tensor));
1286 macro_rules! call {
1287 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1288 let result = <$B>::float_any(unpack!($mode, $variant, Float, tensor, $context), out_dtype);
1289 pack!($mode, $variant, Bool, result)
1290 }};
1291 }
1292 dispatch!(selection, call)
1293 }
1294
1295 fn float_any_dim(tensor: FloatTensor<Self>, dim: usize, out_dtype: BoolDType) -> BoolTensor<Self> {
1296 let selection = Selection::tensors(core::iter::once(&tensor));
1297 macro_rules! call {
1298 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1299 let result = <$B>::float_any_dim(unpack!($mode, $variant, Float, tensor, $context), dim, out_dtype);
1300 pack!($mode, $variant, Bool, result)
1301 }};
1302 }
1303 dispatch!(selection, call)
1304 }
1305
1306 fn float_all(tensor: FloatTensor<Self>, out_dtype: BoolDType) -> BoolTensor<Self> {
1307 let selection = Selection::tensors(core::iter::once(&tensor));
1308 macro_rules! call {
1309 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1310 let result = <$B>::float_all(unpack!($mode, $variant, Float, tensor, $context), out_dtype);
1311 pack!($mode, $variant, Bool, result)
1312 }};
1313 }
1314 dispatch!(selection, call)
1315 }
1316
1317 fn float_all_dim(tensor: FloatTensor<Self>, dim: usize, out_dtype: BoolDType) -> BoolTensor<Self> {
1318 let selection = Selection::tensors(core::iter::once(&tensor));
1319 macro_rules! call {
1320 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1321 let result = <$B>::float_all_dim(unpack!($mode, $variant, Float, tensor, $context), dim, out_dtype);
1322 pack!($mode, $variant, Bool, result)
1323 }};
1324 }
1325 dispatch!(selection, call)
1326 }
1327
1328 fn float_sign(tensor: FloatTensor<Self>) -> FloatTensor<Self> {
1329 let selection = Selection::tensors(core::iter::once(&tensor));
1330 macro_rules! call {
1331 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1332 let result = <$B>::float_sign(unpack!($mode, $variant, Float, tensor, $context));
1333 pack!($mode, $variant, Float, result)
1334 }};
1335 }
1336 dispatch!(selection, call)
1337 }
1338
1339 fn float_expand(tensor: FloatTensor<Self>, shape: Shape) -> FloatTensor<Self> {
1340 let selection = Selection::tensors(core::iter::once(&tensor));
1341 macro_rules! call {
1342 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1343 let result = <$B>::float_expand(unpack!($mode, $variant, Float, tensor, $context), shape);
1344 pack!($mode, $variant, Float, result)
1345 }};
1346 }
1347 dispatch!(selection, call)
1348 }
1349
1350 fn float_sort(tensor: FloatTensor<Self>, dim: usize, descending: bool) -> FloatTensor<Self> {
1351 let selection = Selection::tensors(core::iter::once(&tensor));
1352 macro_rules! call {
1353 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1354 let result = <$B>::float_sort(unpack!($mode, $variant, Float, tensor, $context), dim, descending);
1355 pack!($mode, $variant, Float, result)
1356 }};
1357 }
1358 dispatch!(selection, call)
1359 }
1360
1361 fn float_sort_with_indices( tensor: FloatTensor<Self>, dim: usize, descending: bool, indices_dtype: IntDType, ) -> (FloatTensor<Self>, IntTensor<Self>) {
1362 let selection = Selection::tensors(core::iter::once(&tensor));
1363 macro_rules! call {
1364 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1365 let result = <$B>::float_sort_with_indices(unpack!($mode, $variant, Float, tensor, $context), dim, descending, indices_dtype);
1366 { let (value0, value1) = result; (pack!($mode, $variant, Float, value0), pack!($mode, $variant, Int, value1)) }
1367 }};
1368 }
1369 dispatch!(selection, call)
1370 }
1371
1372 fn float_argsort( tensor: FloatTensor<Self>, dim: usize, descending: bool, out_dtype: IntDType, ) -> IntTensor<Self> {
1373 let selection = Selection::tensors(core::iter::once(&tensor));
1374 macro_rules! call {
1375 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1376 let result = <$B>::float_argsort(unpack!($mode, $variant, Float, tensor, $context), dim, descending, out_dtype);
1377 pack!($mode, $variant, Int, result)
1378 }};
1379 }
1380 dispatch!(selection, call)
1381 }
1382
1383 fn float_grid_sample_2d( tensor: FloatTensor<Self>, grid: FloatTensor<Self>, options: GridSampleOptions, ) -> FloatTensor<Self> {
1384 let selection = Selection::tensors(core::iter::once(&tensor).chain(core::iter::once(&grid)));
1385 macro_rules! call {
1386 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1387 let result = <$B>::float_grid_sample_2d(unpack!($mode, $variant, Float, tensor, $context), unpack!($mode, $variant, Float, grid, $context), options);
1388 pack!($mode, $variant, Float, result)
1389 }};
1390 }
1391 dispatch!(selection, call)
1392 }
1393
1394 fn float_unfold(tensor: FloatTensor<Self>, dim: usize, size: usize, step: usize) -> FloatTensor<Self> {
1395 let selection = Selection::tensors(core::iter::once(&tensor));
1396 macro_rules! call {
1397 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1398 let result = <$B>::float_unfold(unpack!($mode, $variant, Float, tensor, $context), dim, size, step);
1399 pack!($mode, $variant, Float, result)
1400 }};
1401 }
1402 dispatch!(selection, call)
1403 }
1404
1405 fn float_is_nan(tensor: FloatTensor<Self>, out_dtype: BoolDType) -> BoolTensor<Self> {
1406 let selection = Selection::tensors(core::iter::once(&tensor));
1407 macro_rules! call {
1408 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1409 let result = <$B>::float_is_nan(unpack!($mode, $variant, Float, tensor, $context), out_dtype);
1410 pack!($mode, $variant, Bool, result)
1411 }};
1412 }
1413 dispatch!(selection, call)
1414 }
1415
1416 fn float_is_inf(tensor: FloatTensor<Self>, out_dtype: BoolDType) -> BoolTensor<Self> {
1417 let selection = Selection::tensors(core::iter::once(&tensor));
1418 macro_rules! call {
1419 ($B:ty, $variant:ident, $mode:ident, $device:expr, $context:expr) => {{
1420 let result = <$B>::float_is_inf(unpack!($mode, $variant, Float, tensor, $context), out_dtype);
1421 pack!($mode, $variant, Bool, result)
1422 }};
1423 }
1424 dispatch!(selection, call)
1425 }
1426}