burn_tensor/tensor/api/orderable.rs
1use burn_backend::{ElementConversion, Scalar};
2use burn_std::{AsIndex, IndexingUpdateOp};
3
4use crate::check::unwrap_dim_index;
5use crate::kind::Ordered;
6use crate::{Bool, Int, check};
7use crate::{Tensor, check::TensorCheck};
8
9impl<const D: usize, K> Tensor<D, K>
10where
11 K: Ordered,
12{
13 /// Sort the elements by value in ascending order along a given dimension.
14 ///
15 /// This sort is unstable (i.e., may reorder equal elements).
16 ///
17 /// # Arguments
18 ///
19 /// * `dim` - The dimension to sort along.
20 /// Negative dimensions are supported and count from the end.
21 ///
22 /// # Returns
23 ///
24 /// A new tensor with the elements sorted in ascending order along the given dimension.
25 ///
26 /// # Example
27 ///
28 /// ```rust
29 /// use burn_tensor::{Tensor, Shape};
30 ///
31 /// fn example() {
32 /// let device = Default::default();
33 /// let tensor = Tensor::<2>::from_data([[12.0, -2.0, 3.0], [5.0, 3.0, 6.0]], &device);
34 /// let tensor = tensor.sort(0);
35 /// println!("{tensor}");
36 /// // [[5.0, -2.0, 3.0], [12.0, 3.0, 6.0]]
37 /// let tensor = tensor.sort(1);
38 /// println!("{tensor}");
39 /// // [[-2.0, 3.0, 12.0], [3.0, 5.0, 6.0]]
40 /// }
41 /// ```
42 pub fn sort<I: AsIndex>(self, dim: I) -> Self {
43 let dim = unwrap_dim_index(dim.try_dim_index(D), "Sort");
44 Tensor::new(K::sort(self.primitive, dim, /*descending*/ false))
45 }
46
47 /// Sort the elements by value in descending order along a given dimension.
48 ///
49 /// This sort is unstable (i.e., may reorder equal elements).
50 ///
51 /// # Arguments
52 ///
53 /// * `dim` - The dimension to sort along.
54 /// Negative dimensions are supported and count from the end.
55 ///
56 /// # Returns
57 ///
58 /// A new tensor with the elements sorted in descending order along the given dimension.
59 ///
60 /// # Example
61 ///
62 /// ```rust
63 /// use burn_tensor::{Tensor, Shape};
64 ///
65 /// fn example() {
66 /// let device = Default::default();
67 /// let tensor = Tensor::<2>::from_data([[12.0, -2.0, 3.0], [5.0, 3.0, 6.0]], &device);
68 /// let tensor = tensor.sort_descending(0);
69 /// println!("{tensor}");
70 /// // [[12.0, 3.0, 6.0], [5.0, -2.0, 3.0]]
71 /// let tensor = tensor.sort_descending(1);
72 /// println!("{tensor}");
73 /// // [[12.0, 3.0, -2.0], [6.0, 5.0, 3.0]]
74 /// }
75 /// ```
76 pub fn sort_descending<I: AsIndex>(self, dim: I) -> Self {
77 let dim = unwrap_dim_index(dim.try_dim_index(D), "Sort Descending");
78 Tensor::new(K::sort(self.primitive, dim, /*descending*/ true))
79 }
80
81 /// Sort the elements by value in ascending order along a given dimension.
82 /// Also returns the indices.
83 ///
84 /// This sort is unstable (i.e., may reorder equal elements).
85 ///
86 /// # Arguments
87 ///
88 /// * `dim` - The dimension to sort along.
89 /// Negative dimensions are supported and count from the end.
90 ///
91 /// # Returns
92 ///
93 /// A tuple containing the sorted tensor and the indices tensor.
94 ///
95 /// # Example
96 ///
97 /// ```rust
98 /// use burn_tensor::{Tensor, Shape};
99 ///
100 /// fn example() {
101 /// let device = Default::default();
102 /// let tensor = Tensor::<2>::from_data([[12.0, -2.0, 3.0], [5.0, 3.0, 6.0]], &device);
103 /// let (tensor, indices) = tensor.sort_with_indices(0);
104 /// println!("{tensor}");
105 /// // [[5.0, -2.0, 3.0], [12.0, 3.0, 6.0]]
106 /// println!("{}", indices);
107 /// // [[1, 0, 0], [0, 1, 1]]
108 /// }
109 /// ```
110 pub fn sort_with_indices<I: AsIndex>(self, dim: I) -> (Self, Tensor<D, Int>) {
111 let dim = unwrap_dim_index(dim.try_dim_index(D), "Sort With Indices");
112 let (values, indices) =
113 K::sort_with_indices(self.primitive, dim, /*descending*/ false);
114 (Tensor::new(values), Tensor::new(indices))
115 }
116
117 /// Sort the elements by value in descending order along a given dimension.
118 /// Also returns the indices.
119 ///
120 /// This sort is unstable (i.e., may reorder equal elements).
121 ///
122 /// # Arguments
123 ///
124 /// * `dim` - The dimension to sort along.
125 /// Negative dimensions are supported and count from the end.
126 ///
127 /// # Example
128 ///
129 /// ```rust
130 /// use burn_tensor::{Tensor, Shape};
131 ///
132 /// fn example() {
133 /// let device = Default::default();
134 /// let tensor = Tensor::<2>::from_data([[12.0, -2.0, 3.0], [5.0, 3.0, 6.0]], &device);
135 /// let (tensor, indices) = tensor.sort_descending_with_indices(0);
136 /// println!("{tensor}");
137 /// // [[12.0, 3.0, 6.0], [5.0, -2.0, 3.0]]
138 /// println!("{}", indices);
139 /// // [[0, 1, 1], [1, 0, 0]]
140 /// }
141 /// ```
142 pub fn sort_descending_with_indices<I: AsIndex>(self, dim: I) -> (Self, Tensor<D, Int>) {
143 let dim = unwrap_dim_index(dim.try_dim_index(D), "Sort Descending With Indices");
144 let (values, indices) = K::sort_with_indices(self.primitive, dim, /*descending*/ true);
145 (Tensor::new(values), Tensor::new(indices))
146 }
147
148 /// Returns the indices that sort the elements by value in ascending order along a given dimension.
149 ///
150 /// This sort is unstable (i.e., may reorder equal elements).
151 ///
152 /// # Arguments
153 ///
154 /// * `dim` - The dimension to sort along.
155 /// Negative dimensions are supported and count from the end.
156 ///
157 /// # Example
158 ///
159 /// ```rust
160 /// use burn_tensor::{Tensor, Shape};
161 ///
162 /// fn example() {
163 /// let device = Default::default();
164 /// let tensor = Tensor::<2>::from_data([[12.0, -2.0, 3.0], [5.0, 3.0, 6.0]], &device);
165 /// let tensor = tensor.argsort(0);
166 /// println!("{tensor}");
167 /// // [[1, 0, 0], [0, 1, 1]]
168 /// }
169 /// ```
170 pub fn argsort<I: AsIndex>(self, dim: I) -> Tensor<D, Int> {
171 let dim = unwrap_dim_index(dim.try_dim_index(D), "Argsort");
172 Tensor::new(K::argsort(self.primitive, dim, /*descending*/ false))
173 }
174
175 /// Returns the indices that sort the elements by value in descending order along a given dimension.
176 ///
177 /// This sort is unstable (i.e., may reorder equal elements).
178 ///
179 /// # Arguments
180 ///
181 /// * `dim` - The dimension to sort along.
182 /// Negative dimensions are supported and count from the end.
183 ///
184 /// # Example
185 ///
186 /// ```rust
187 /// use burn_tensor::{Tensor, Shape};
188 ///
189 /// fn example() {
190 /// let device = Default::default();
191 /// let tensor = Tensor::<2>::from_data([[12.0, -2.0, 3.0], [5.0, 3.0, 6.0]], &device);
192 /// let tensor = tensor.argsort_descending(0);
193 /// println!("{tensor}");
194 /// // [[0, 1, 1], [1, 0, 0]]
195 /// let tensor = tensor.argsort_descending(1);
196 /// println!("{tensor}");
197 /// // [[0, 2, 1], [2, 0, 1]]
198 /// }
199 /// ```
200 pub fn argsort_descending<I: AsIndex>(self, dim: I) -> Tensor<D, Int> {
201 let dim = unwrap_dim_index(dim.try_dim_index(D), "Argsort Descending");
202 Tensor::new(K::argsort(self.primitive, dim, /*descending*/ true))
203 }
204
205 /// Returns the `k` largest elements of the given input tensor along a given dimension.
206 ///
207 /// # Arguments
208 ///
209 /// * `k` - The number of elements to return.
210 /// * `dim` - The dimension to sort along.
211 /// Negative dimensions are supported and count from the end.
212 ///
213 /// # Returns
214 ///
215 /// A new tensor with the `k` largest elements along the given dimension.
216 ///
217 /// # Example
218 ///
219 /// ```rust
220 /// use burn_tensor::{Tensor, Shape};
221 ///
222 /// fn example() {
223 /// let device = Default::default();
224 /// let tensor = Tensor::<2>::from_data([[12.0, -2.0, 3.0], [5.0, 3.0, 6.0]], &device);
225 /// let tensor = tensor.topk(2, 0);
226 /// println!("{tensor}");
227 /// // [[12.0, 3.0, 6.0], [5.0, -2.0, 3.0]]
228 /// let tensor = tensor.topk(1, 1);
229 /// println!("{tensor}");
230 /// // [[12.0], [6.0]]
231 /// }
232 /// ```
233 pub fn topk<I: AsIndex>(self, k: usize, dim: I) -> Self {
234 let dim = unwrap_dim_index(dim.try_dim_index(D), "Top K");
235 assert!(self.shape()[dim] > k);
236 Tensor::new(K::topk(self.primitive, dim, k))
237 }
238
239 /// Returns the `k` largest elements of the given input tensor along a given dimension.
240 /// Also returns the indices.
241 ///
242 /// # Arguments
243 ///
244 /// * `k` - The number of elements to return.
245 /// * `dim` - The dimension to sort along.
246 /// Negative dimensions are supported and count from the end.
247 ///
248 /// # Example
249 ///
250 /// ```rust
251 /// use burn_tensor::{Tensor, Shape};
252 ///
253 /// fn example() {
254 /// let device = Default::default();
255 /// let tensor = Tensor::<2>::from_data([[12.0, -2.0, 3.0], [5.0, 3.0, 6.0]], &device);
256 /// let (tensor, indices) = tensor.topk_with_indices(2, 0);
257 /// println!("{tensor}");
258 /// // [[12.0, 3.0, 6.0], [5.0, -2.0, 3.0]]
259 /// println!("{}", indices);
260 /// // [[0, 1, 1], [1, 0, 0]]
261 /// let (tensor, indices) = tensor.topk_with_indices(1, 1);
262 /// println!("{tensor}");
263 /// // [[12.0], [6.0]]
264 /// println!("{indices}");
265 /// // [[0], [2]]
266 /// }
267 /// ```
268 pub fn topk_with_indices<I: AsIndex>(self, k: usize, dim: I) -> (Self, Tensor<D, Int>) {
269 let dim = unwrap_dim_index(dim.try_dim_index(D), "Top K With Indices");
270 let (values, indices) = K::topk_with_indices(self.primitive, dim, k);
271 (Tensor::new(values), Tensor::new(indices))
272 }
273
274 /// Create a one hot tensor.
275 ///
276 /// # Example
277 ///
278 /// ```rust
279 /// use burn_tensor::Tensor;
280 ///
281 /// fn example(){
282 /// let device = Default::default();
283 /// let indices: Tensor<1> = Tensor::from_floats([0.0, 1.0, 2.0, 3.0], &device);
284 /// let one_hot: Tensor<2> = indices.one_hot(4);
285 /// println!("{}", one_hot.to_data());
286 /// // [[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 1.0]]
287 /// }
288 /// ```
289 pub fn one_hot<const D2: usize>(self, num_classes: usize) -> Tensor<D2, K> {
290 check!(TensorCheck::one_hot_tensor(self.clone(), num_classes));
291 self.one_hot_fill(num_classes, 1.0, 0.0, -1)
292 }
293
294 /// Create a one-hot encoded tensor with configurable `num_classes`, `on_value`, `off_value`, and `axis` including high-ranked tensors.
295 ///
296 /// # Arguments
297 ///
298 /// * `num_classes`: The number of classes for the one-hot encoding, which defines the size of the one-hot dimension.
299 /// * `on_value`: The value to assign for active positions (corresponding to indices).
300 /// * `off_value`: The value to assign for inactive positions.
301 /// * `axis`: The axis along which the one-hot dimension is added.
302 /// Negative dimensions are supported and count from the end.
303 ///
304 /// # Returns
305 ///
306 /// A tensor with one additional dimension for the one-hot encoding, where active positions are filled with `on_value` and others with `off_value`.
307 ///
308 /// # Example
309 /// ```rust
310 /// use burn_tensor::{Tensor, Float};
311 /// fn example() {
312 /// let device = Default::default();
313 /// let indices: Tensor<2, Float> = Tensor::from_floats([[0., 2.], [1., -1.]], &device);
314 /// // One-hot encoding
315 /// let tensor: Tensor<3, Float> = indices.one_hot_fill(3, 5.0.into(), 0.0.into(), -1);
316 /// println!("{tensor}");
317 /// // [[[5.0, 0.0, 0.0],
318 /// // [0.0, 0.0, 5.0]],
319 /// // [[0.0, 5.0, 0.0],
320 /// // [0.0, 0.0, 5.0]]]
321 /// }
322 /// ```
323 pub fn one_hot_fill<const D2: usize>(
324 self,
325 num_classes: usize,
326 on_value: f32,
327 off_value: f32,
328 axis: impl AsIndex,
329 ) -> Tensor<D2, K> {
330 check!(TensorCheck::one_hot_tensor_rank::<D, D2>());
331 let axis = unwrap_dim_index(axis.try_dim_index(D + 1), "One Hot");
332
333 // Initialize shape from the current tensor dimensions and prepare for modification
334 let mut shape = self.shape();
335 let device = self.device();
336
337 // Convert the input tensor to integer indices
338 let indices: Tensor<D, Int> = Tensor::from_data(self.to_data().convert::<i64>(), &device);
339 // Insert the new dimension for the one-hot representation
340 shape.insert(axis, num_classes);
341 // Adjust indices to valid range and handle invalid indices
342 let adjusted_indices = indices
343 .clone()
344 .mask_fill(self.clone().lower_scalar(0), num_classes as i64) // Handle negative indices
345 .add(indices.clone().mask_fill(self.clone().greater_scalar(0), 0)); // Handle positive indices
346 // Unsqueeze the indices tensor along the specified axis
347 let indices_unsqueezed: Tensor<D2, Int> = adjusted_indices.unsqueeze_dim(axis);
348
349 // Initialize the output tensor with the off_value
350 let output = Tensor::full(shape.clone(), off_value, &device);
351
352 // Prepare scatter tensor for on_value and off_value adjustments
353 let scatter_on_values = Tensor::full(indices_unsqueezed.shape(), on_value, &device)
354 - Tensor::full(indices_unsqueezed.shape(), off_value, &self.device());
355
356 // Scatter on_value at the appropriate indices to create the one-hot representation
357 output.scatter(
358 axis,
359 indices_unsqueezed,
360 scatter_on_values,
361 IndexingUpdateOp::Add,
362 )
363 }
364
365 /// Applies element wise greater comparison and returns a boolean tensor.
366 ///
367 /// # Panics
368 ///
369 /// If the two tensors don't have the same shape.
370 ///
371 /// # Example
372 ///
373 /// ```rust
374 /// use burn_tensor::{Tensor, Shape};
375 ///
376 /// fn example() {
377 /// let device = Default::default();
378 /// let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
379 /// let tensor2 = Tensor::<2>::from_data([[1.0, 3.0, 4.0], [1.0, 2.0, 3.0]], &device);
380 /// let tensor = tensor1.greater(tensor2);
381 /// println!("{tensor}");
382 /// // [[false, false, false], [true, true, true]]
383 /// }
384 /// ```
385 pub fn greater(self, other: Self) -> Tensor<D, Bool> {
386 check!(TensorCheck::binary_ops_ew("Greater", &self, &other));
387 Tensor::new(K::greater(self.primitive, other.primitive))
388 }
389
390 /// Applies element wise greater-equal comparison and returns a boolean tensor.
391 ///
392 /// # Panics
393 ///
394 /// If the two tensors don't have the same shape.
395 ///
396 /// # Example
397 ///
398 /// ```rust
399 /// use burn_tensor::{Tensor, Shape};
400 ///
401 /// fn example() {
402 /// let device = Default::default();
403 /// let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
404 /// let tensor2 = Tensor::<2>::from_data([[1.0, 3.0, 4.0], [1.0, 2.0, 3.0]], &device);
405 /// let tensor = tensor1.greater_equal(tensor2);
406 /// println!("{tensor}");
407 /// // [[true, false, false], [true, true, true]]
408 /// }
409 /// ```
410 pub fn greater_equal(self, other: Self) -> Tensor<D, Bool> {
411 check!(TensorCheck::binary_ops_ew("Greater_equal", &self, &other));
412 Tensor::new(K::greater_equal(self.primitive, other.primitive))
413 }
414
415 /// Applies element wise lower comparison and returns a boolean tensor.
416 ///
417 /// # Panics
418 ///
419 /// If the two tensors don't have the same shape.
420 ///
421 /// # Example
422 ///
423 /// ```rust
424 /// use burn_tensor::{Tensor, Shape};
425 ///
426 /// fn example() {
427 /// let device = Default::default();
428 /// let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
429 /// let tensor2 = Tensor::<2>::from_data([[1.0, 3.0, 4.0], [1.0, 2.0, 3.0]], &device);
430 /// let tensor = tensor1.lower(tensor2);
431 /// println!("{tensor}");
432 /// // [[false, true, true], [false, false, false]]
433 /// }
434 /// ```
435 pub fn lower(self, other: Self) -> Tensor<D, Bool> {
436 check!(TensorCheck::binary_ops_ew("Lower", &self, &other));
437 Tensor::new(K::lower(self.primitive, other.primitive))
438 }
439
440 /// Applies element wise lower-equal comparison and returns a boolean tensor.
441 ///
442 /// # Panics
443 ///
444 /// If the two tensors don't have the same shape.
445 ///
446 /// # Example
447 ///
448 /// ```rust
449 /// use burn_tensor::{Tensor, Shape};
450 ///
451 /// fn example() {
452 /// let device = Default::default();
453 /// let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
454 /// let tensor2 = Tensor::<2>::from_data([[1.0, 3.0, 4.0], [1.0, 2.0, 3.0]], &device);
455 /// let tensor = tensor1.lower_equal(tensor2);
456 /// println!("{tensor}");
457 /// // [[true, true, true], [false, false, false]]
458 /// }
459 /// ```
460 pub fn lower_equal(self, other: Self) -> Tensor<D, Bool> {
461 check!(TensorCheck::binary_ops_ew("Lower_equal", &self, &other));
462 Tensor::new(K::lower_equal(self.primitive, other.primitive))
463 }
464
465 /// Applies greater than `other` comparison and returns a boolean tensor.
466 ///
467 /// # Arguments
468 ///
469 /// * `other` - The scalar to compare.
470 ///
471 /// # Example
472 ///
473 /// ```rust
474 /// use burn_tensor::{Tensor, Shape};
475 ///
476 /// fn example() {
477 /// let device = Default::default();
478 /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
479 /// let tensor = tensor.greater_scalar(3.0);
480 /// println!("{tensor}");
481 /// // [[false, false, true], [true, true, true]]
482 /// }
483 /// ```
484 pub fn greater_scalar<E: ElementConversion>(self, other: E) -> Tensor<D, Bool> {
485 let other = Scalar::new(other, &self.dtype());
486 Tensor::new(K::greater_scalar(self.primitive, other))
487 }
488
489 /// Applies greater-equal than `other` comparison and returns a boolean tensor.
490 ///
491 /// # Arguments
492 ///
493 /// * `other` - The scalar to compare.
494 ///
495 /// # Example
496 ///
497 /// ```rust
498 /// use burn_tensor::{Tensor, Shape};
499 ///
500 /// fn example() {
501 /// let device = Default::default();
502 /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
503 /// let tensor = tensor.greater_equal_scalar(3.0);
504 /// println!("{tensor}");
505 /// // [[false, false, true], [true, true, true]]
506 /// }
507 /// ```
508 pub fn greater_equal_scalar<E: ElementConversion>(self, other: E) -> Tensor<D, Bool> {
509 let other = Scalar::new(other, &self.dtype());
510 Tensor::new(K::greater_equal_scalar(self.primitive, other))
511 }
512
513 /// Applies lower than `other` comparison and returns a boolean tensor.
514 ///
515 /// # Arguments
516 ///
517 /// * `other` - The scalar to compare.
518 ///
519 /// # Example
520 ///
521 /// ```rust
522 /// use burn_tensor::{Tensor, Shape};
523 ///
524 /// fn example() {
525 /// let device = Default::default();
526 /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
527 /// let tensor = tensor.lower_scalar(3.0);
528 /// println!("{tensor}");
529 /// // [[true, true, false], [false, false, false]]
530 /// }
531 /// ```
532 pub fn lower_scalar<E: ElementConversion>(self, other: E) -> Tensor<D, Bool> {
533 let other = Scalar::new(other, &self.dtype());
534 Tensor::new(K::lower_scalar(self.primitive, other))
535 }
536
537 /// Applies lower-equal than `other` comparison and returns a boolean tensor.
538 ///
539 /// # Arguments
540 ///
541 /// * `other` - The scalar to compare.
542 ///
543 /// # Example
544 ///
545 /// ```rust
546 /// use burn_tensor::{Tensor, Shape};
547 ///
548 /// fn example() {
549 /// let device = Default::default();
550 /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
551 /// let tensor = tensor.lower_equal_scalar(3.0);
552 /// println!("{tensor}");
553 /// // [[true, true, true], [false, false, false]]
554 /// }
555 /// ```
556 pub fn lower_equal_scalar<E: ElementConversion>(self, other: E) -> Tensor<D, Bool> {
557 let other = Scalar::new(other, &self.dtype());
558 Tensor::new(K::lower_equal_scalar(self.primitive, other))
559 }
560
561 /// Alias for [greater_scalar](Self::greater_scalar).
562 pub fn greater_elem<E: ElementConversion>(self, other: E) -> Tensor<D, Bool> {
563 self.greater_scalar(other)
564 }
565
566 /// Alias for [greater_equal_scalar](Self::greater_equal_scalar).
567 pub fn greater_equal_elem<E: ElementConversion>(self, other: E) -> Tensor<D, Bool> {
568 self.greater_equal_scalar(other)
569 }
570
571 /// Alias for [lower_scalar](Self::lower_scalar).
572 pub fn lower_elem<E: ElementConversion>(self, other: E) -> Tensor<D, Bool> {
573 self.lower_scalar(other)
574 }
575
576 /// Alias for [lower_equal_scalar](Self::lower_equal_scalar).
577 pub fn lower_equal_elem<E: ElementConversion>(self, other: E) -> Tensor<D, Bool> {
578 self.lower_equal_scalar(other)
579 }
580
581 /// Applies the argmax function along the given dimension and returns an integer tensor.
582 ///
583 /// # Arguments
584 ///
585 /// * `dim` - The dimension along which to find the maximum value.
586 /// Negative dimensions are supported and count from the end.
587 ///
588 /// # Example
589 ///
590 /// ```rust
591 /// use burn_tensor::{Tensor, Shape};
592 ///
593 /// fn example() {
594 /// let device = Default::default();
595 /// let tensor = Tensor::<3>::ones(Shape::new([2, 3, 3]), &device);
596 /// let tensor = tensor.argmax(1);
597 /// println!("{:?}", tensor.shape());
598 /// // Shape { dims: [2, 1, 3] }
599 /// }
600 /// ```
601 pub fn argmax(self, dim: impl AsIndex) -> Tensor<D, Int> {
602 let dim = unwrap_dim_index(dim.try_dim_index(D), "Argmax");
603 Tensor::new(K::argmax(self.primitive, dim))
604 }
605
606 /// Applies the argtopk function along the given dimension and returns an integer tensor.
607 ///
608 /// # Arguments
609 ///
610 /// * `k` - The number of indices to return.
611 /// * `dim` - The dimension along which to find the largest values.
612 /// Negative dimensions are supported and count from the end.
613 ///
614 /// # Example
615 ///
616 /// ```rust
617 /// use burn_tensor::{Tensor, Shape};
618 ///
619 /// fn example() {
620 /// let device = Default::default();
621 /// let tensor = Tensor::<3>::ones(Shape::new([2, 3, 3]), &device);
622 /// let tensor = tensor.argtopk(1, 2);
623 /// println!("{:?}", tensor.shape());
624 /// }
625 /// ```
626 pub fn argtopk(self, k: usize, dim: impl AsIndex) -> Tensor<D, Int> {
627 let dim = unwrap_dim_index(dim.try_dim_index(D), "Argtopk");
628 assert!(self.shape()[dim] > k);
629 Tensor::new(K::argtopk(self.primitive, dim, k))
630 }
631
632 /// Find the maximum value.
633 ///
634 /// # Example
635 ///
636 /// ```rust
637 /// use burn_tensor::{Tensor, Shape};
638 ///
639 /// fn example() {
640 /// let device = Default::default();
641 /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
642 /// let tensor = tensor.max();
643 /// println!("{tensor}");
644 /// // [9.0]
645 /// }
646 /// ```
647 pub fn max(self) -> Tensor<1, K> {
648 Tensor::new(K::max(self.primitive))
649 }
650
651 /// Find the maximum value along the given dimension.
652 ///
653 /// Also returns the indices.
654 ///
655 /// # Example
656 ///
657 /// ```rust
658 /// use burn_tensor::{Tensor, Shape};
659 ///
660 /// fn example() {
661 /// let device = Default::default();
662 /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
663 /// let (tensor, index) = tensor.max_dim_with_indices(0);
664 /// // [[5.0, 9.0, 6.0]]
665 /// println!("{tensor}");
666 /// // [[1, 1, 1]]
667 /// println!("{index}");
668 /// }
669 /// ```
670 pub fn max_dim_with_indices<I: AsIndex>(self, dim: I) -> (Self, Tensor<D, Int>) {
671 let dim = unwrap_dim_index(dim.try_dim_index(D), "Max Dim With Indices");
672
673 let (tensor, index) = K::max_dim_with_indices(self.primitive, dim);
674
675 let tensor = Tensor::new(tensor);
676 let index = Tensor::new(index);
677
678 (tensor, index)
679 }
680
681 /// Find the maximum absolute value.
682 ///
683 /// # Example
684 ///
685 /// ```rust
686 /// use burn_tensor::{Tensor, Shape};
687 ///
688 /// fn example() {
689 /// let device = Default::default();
690 /// let tensor = Tensor::<2>::from_data([[1.0, -7.0, 3.0], [5.0, -1.0, 6.0]], &device);
691 /// let tensor = tensor.max_abs();
692 /// println!("{tensor}");
693 /// // [7.0]
694 /// }
695 /// ```
696 pub fn max_abs(self) -> Tensor<1, K> {
697 Tensor::new(K::max_abs(self.primitive))
698 }
699
700 /// Finds the maximum pair wise values with another tensor.
701 ///
702 /// # Arguments
703 ///
704 /// * `other` - Other tensor to find maximum elements with
705 ///
706 /// # Returns
707 ///
708 /// A tensor with the same shape as the input tensors containing the maximum value found
709 /// in the input tensors.
710 ///
711 /// # Example
712 ///
713 /// ```rust
714 /// use burn_tensor::{Tensor, Shape};
715 ///
716 /// fn example() {
717 /// let device = Default::default();
718 /// let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
719 /// let tensor2 = Tensor::<2>::from_data([[2.0, 3.0, 4.0], [1.0, 2.0, 3.0]], &device);
720 /// let tensor = tensor1.max_pair(tensor2);
721 /// println!("{tensor}");
722 /// // [[2.0, 3.0, 4.0], [5.0, 9.0, 6.0]]
723 /// }
724 /// ```
725 pub fn max_pair(self, other: Self) -> Self {
726 let mask = self.clone().lower(other.clone());
727 self.mask_where(mask, other)
728 }
729
730 /// Find the maximum absolute value along the given dimension.
731 ///
732 /// # Arguments
733 ///
734 /// * `dim` - The dimension or axis along which to aggregate the elements,
735 /// supports negative indexing.
736 ///
737 /// # Returns
738 ///
739 /// The returned tensor will have the same rank,
740 /// but the aggregated dimension will have size 1.
741 ///
742 /// # Example
743 ///
744 /// ```rust
745 /// use burn_tensor::{Tensor, Shape};
746 ///
747 /// fn example() {
748 /// let device = Default::default();
749 /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
750 /// let tensor = tensor.max_dim(0);
751 /// println!("{tensor}");
752 /// // [[5.0, 9.0, 6.0]]
753 /// }
754 /// ```
755 pub fn max_abs_dim<I: AsIndex>(self, dim: I) -> Self {
756 let dim = unwrap_dim_index(dim.try_dim_index(D), "Max Abs Dim");
757
758 Tensor::new(K::max_abs_dim(self.primitive, dim))
759 }
760
761 /// Find the maximum absolute value along the given dimensions.
762 ///
763 /// # Arguments
764 ///
765 /// * `dims` - The dimensions or axes along which to aggregate the elements,
766 /// supports negative indexing.
767 ///
768 /// # Returns
769 ///
770 /// The returned tensor will have the same rank,
771 /// but the aggregated dimensions will have size 1.
772 ///
773 /// # Example
774 ///
775 /// ```rust
776 /// use burn_tensor::{Tensor, Shape};
777 ///
778 /// fn example() {
779 /// let device = Default::default();
780 /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
781 /// let tensor = tensor.max_abs_dims(&[0, 1]);
782 /// println!("{tensor}");
783 /// // [[9.0]]
784 /// }
785 /// ```
786 pub fn max_abs_dims<I: AsIndex>(self, dims: &[I]) -> Self {
787 dims.iter()
788 .fold(self, |tensor, &dim| tensor.max_abs_dim(dim))
789 }
790
791 /// Applies the argmin function along the given dimension and returns an integer tensor.
792 ///
793 /// # Arguments
794 ///
795 /// * `dim` - The dimension along which to find the minimum value.
796 /// Negative dimensions are supported and count from the end.
797 ///
798 /// # Example
799 ///
800 /// ```rust
801 /// use burn_tensor::{Tensor, Shape};
802 ///
803 /// fn example() {
804 /// let device = Default::default();
805 /// let tensor = Tensor::<3>::ones(Shape::new([2, 3, 3]), &device);
806 /// let tensor = tensor.argmin(1);
807 /// println!("{:?}", tensor.shape());
808 /// // Shape { dims: [2, 1, 3] }
809 /// }
810 /// ```
811 pub fn argmin(self, dim: impl AsIndex) -> Tensor<D, Int> {
812 let dim = unwrap_dim_index(dim.try_dim_index(D), "Argmin");
813 Tensor::new(K::argmin(self.primitive, dim))
814 }
815
816 /// Find the minimum value.
817 ///
818 /// # Example
819 ///
820 /// ```rust
821 /// use burn_tensor::{Tensor, Shape};
822 ///
823 /// fn example() {
824 /// let device = Default::default();
825 /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
826 /// let tensor = tensor.min();
827 /// println!("{tensor}");
828 /// // [-2.0]
829 /// }
830 /// ```
831 pub fn min(self) -> Tensor<1, K> {
832 Tensor::new(K::min(self.primitive))
833 }
834
835 /// Find the minimum value along the given dimension.
836 ///
837 /// # Arguments
838 ///
839 /// * `dim` - The dimension or axis along which to aggregate the elements;
840 /// supports negative indexing.
841 ///
842 /// # Returns
843 ///
844 /// The returned tensor will have the same rank,
845 /// but the aggregated dimension will have size 1.
846 ///
847 /// # Example
848 ///
849 /// ```rust
850 /// use burn_tensor::{Tensor, Shape};
851 ///
852 /// fn example() {
853 /// let device = Default::default();
854 /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
855 /// let tensor = tensor.min_dim(0);
856 /// println!("{tensor}");
857 /// // [[1.0, -2.0, 3.0]]
858 /// }
859 /// ```
860 pub fn min_dim<I: AsIndex>(self, dim: I) -> Self {
861 let dim = unwrap_dim_index(dim.try_dim_index(D), "Min Dim");
862 Tensor::new(K::min_dim(self.primitive, dim))
863 }
864
865 /// Find the minimum value along the given dimensions.
866 ///
867 /// # Arguments
868 ///
869 /// * `dims` - The dimensions or axes along which to aggregate the elements;
870 /// supports negative indexing.
871 ///
872 /// # Returns
873 ///
874 /// The returned tensor will have the same rank,
875 /// but the aggregated dimensions will have size 1.
876 ///
877 /// # Example
878 ///
879 /// ```rust
880 /// use burn_tensor::{Tensor, Shape};
881 ///
882 /// fn example() {
883 /// let device = Default::default();
884 /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
885 /// let tensor = tensor.min_dims(&[0, 1]);
886 /// println!("{tensor}");
887 /// // [[-2.0]]
888 /// }
889 /// ```
890 pub fn min_dims<I: AsIndex>(self, dims: &[I]) -> Self {
891 dims.iter().fold(self, |tensor, &dim| tensor.min_dim(dim))
892 }
893
894 /// Find the minimum value along the given dimension.
895 ///
896 /// Also returns the indices.
897 ///
898 /// # Example
899 ///
900 /// ```rust
901 /// use burn_tensor::{Tensor, Shape};
902 ///
903 /// fn example() {
904 /// let device = Default::default();
905 /// let tensor = Tensor::<2>::from_data([[7.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
906 /// let (tensor, index) = tensor.min_dim_with_indices(0);
907 /// println!("{tensor}");
908 /// // [[5.0, -2.0, 3.0]]
909 /// println!("{}", index);
910 /// // [[1, 0, 0]]
911 /// }
912 /// ```
913 pub fn min_dim_with_indices<I: AsIndex>(self, dim: I) -> (Self, Tensor<D, Int>) {
914 let dim = unwrap_dim_index(dim.try_dim_index(D), "Min Dim With Indices");
915
916 let (tensor, index) = K::min_dim_with_indices(self.primitive, dim);
917
918 let tensor = Tensor::new(tensor);
919 let index = Tensor::new(index);
920
921 (tensor, index)
922 }
923
924 /// Finds the minimum pair wise values with another tensor.
925 ///
926 /// # Arguments
927 ///
928 /// * `other` - Other tensor to find minimum elements with
929 ///
930 /// # Returns
931 ///
932 /// A tensor with the same shape as the input tensors containing the minimum value found
933 /// between each element of the two source tensors.
934 ///
935 /// # Example
936 ///
937 /// ```rust
938 /// use burn_tensor::{Tensor, Shape};
939 ///
940 /// fn example() {
941 /// let device = Default::default();
942 /// let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
943 /// let tensor2 = Tensor::<2>::from_data([[2.0, 3.0, 4.0], [1.0, 2.0, 3.0]], &device);
944 /// let tensor = tensor1.min_pair(tensor2);
945 /// println!("{tensor}");
946 /// // [[1.0, -2.0, 3.0], [1.0, 2.0, 3.0]]
947 /// }
948 pub fn min_pair(self, other: Self) -> Self {
949 let mask = other.clone().lower(self.clone());
950 self.mask_where(mask, other)
951 }
952
953 /// Clamp element wise between the given min and max values.
954 ///
955 /// # Arguments
956 ///
957 /// * `min` - The minimum value.
958 /// * `max` - The maximum value.
959 ///
960 /// # Returns
961 ///
962 /// A new tensor with the values clamped between the given min and max values.
963 ///
964 /// # Example
965 ///
966 /// ```rust
967 /// use burn_tensor::{Int, Tensor};
968 ///
969 /// fn example() {
970 /// let device = Default::default();
971 /// let tensor = Tensor::<2, Int>::from_ints(
972 /// [
973 /// [1, 2, 3],
974 /// [4, 5, 6],
975 /// [7, 8, 9]
976 /// ],
977 /// &device);
978 /// let tensor = tensor.clamp(2, 6);
979 /// println!("{tensor}");
980 /// // [[2, 2, 3], [4, 5, 6], [6, 6, 6]]
981 /// }
982 /// ```
983 pub fn clamp<E: ElementConversion>(self, min: E, max: E) -> Self {
984 let dtype = self.dtype();
985 Self::new(K::clamp(
986 self.primitive,
987 Scalar::new(min, &dtype),
988 Scalar::new(max, &dtype),
989 ))
990 }
991
992 /// Clamp element wise under a minimum value.
993 ///
994 /// # Arguments
995 ///
996 /// * `tensor` - The tensor to clamp.
997 /// * `min` - The minimum value.
998 ///
999 /// # Returns
1000 ///
1001 /// A new tensor with the values clamped under the given min value.
1002 ///
1003 /// # Example
1004 ///
1005 /// ```rust
1006 /// use burn_tensor::{Int, Tensor};
1007 ///
1008 /// fn example() {
1009 /// let device = Default::default();
1010 /// let tensor = Tensor::<2, Int>::from_ints(
1011 /// [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
1012 /// &device);
1013 /// let tensor = tensor.clamp_min(4);
1014 /// println!("{tensor}");
1015 /// // [[4, 4, 4], [4, 5, 6], [7, 8, 9]]
1016 /// }
1017 /// ```
1018 pub fn clamp_min<E: ElementConversion>(self, min: E) -> Self {
1019 let min = Scalar::new(min, &self.dtype());
1020 Self::new(K::clamp_min(self.primitive, min))
1021 }
1022
1023 /// Clamp element wise over a maximum value.
1024 ///
1025 /// # Arguments
1026 ///
1027 /// * `tensor` - The tensor to clamp.
1028 /// * `max` - The maximum value.
1029 ///
1030 /// # Returns
1031 ///
1032 /// A new tensor with the values clamped over the given max value.
1033 ///
1034 /// # Example
1035 ///
1036 /// ```rust
1037 /// use burn_tensor::{Int, Tensor};
1038 ///
1039 /// fn example() {
1040 /// let device = Default::default();
1041 /// let tensor = Tensor::<2, Int>::from_ints(
1042 /// [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
1043 /// &device);
1044 /// let tensor = tensor.clamp_max(5);
1045 /// println!("{tensor}");
1046 /// // [[1, 2, 3], [4, 5, 5], [5, 5, 5]]
1047 /// }
1048 /// ```
1049 pub fn clamp_max<E: ElementConversion>(self, max: E) -> Self {
1050 let max = Scalar::new(max, &self.dtype());
1051 Self::new(K::clamp_max(self.primitive, max))
1052 }
1053
1054 /// Computes the cumulative minimum of elements along the given *dimension* or *axis*.
1055 ///
1056 /// # Arguments
1057 ///
1058 /// * `dim` - The dimension or axis along which to compute the cumulative minimum.
1059 /// Negative dimensions are supported and count from the end.
1060 ///
1061 /// # Example
1062 ///
1063 /// ```rust
1064 /// use burn_tensor::{Tensor, Shape};
1065 ///
1066 /// fn example() {
1067 /// let device = Default::default();
1068 /// let tensor = Tensor::<2>::from_data([[3.0, 5.0, 2.0], [4.0, 1.0, 6.0]], &device);
1069 /// let result = tensor.clone().cummin(0);
1070 /// println!("{result}");
1071 /// // [[3.0, 5.0, 2.0], [3.0, 1.0, 2.0]]
1072 /// let result = tensor.cummin(1);
1073 /// println!("{result}");
1074 /// // [[3.0, 3.0, 2.0], [4.0, 1.0, 1.0]]
1075 /// }
1076 /// ```
1077 pub fn cummin<I: AsIndex>(self, dim: I) -> Self {
1078 let dim = unwrap_dim_index(dim.try_dim_index(D), "Cummin");
1079 Self::new(K::cummin(self.primitive, dim))
1080 }
1081
1082 /// Computes the cumulative maximum of elements along the given *dimension* or *axis*.
1083 ///
1084 /// # Arguments
1085 ///
1086 /// * `dim` - The dimension or axis along which to compute the cumulative maximum.
1087 /// Negative dimensions are supported and count from the end.
1088 ///
1089 /// # Example
1090 ///
1091 /// ```rust
1092 /// use burn_tensor::{Tensor, Shape};
1093 ///
1094 /// fn example() {
1095 /// let device = Default::default();
1096 /// let tensor = Tensor::<2>::from_data([[3.0, 1.0, 2.0], [4.0, 5.0, 2.0]], &device);
1097 /// let result = tensor.clone().cummax(0);
1098 /// println!("{result}");
1099 /// // [[3.0, 1.0, 2.0], [4.0, 5.0, 2.0]]
1100 /// let result = tensor.cummax(1);
1101 /// println!("{result}");
1102 /// // [[3.0, 3.0, 3.0], [4.0, 5.0, 5.0]]
1103 /// }
1104 /// ```
1105 pub fn cummax<I: AsIndex>(self, dim: I) -> Self {
1106 let dim = unwrap_dim_index(dim.try_dim_index(D), "Cummax");
1107 Self::new(K::cummax(self.primitive, dim))
1108 }
1109 /// Find the maximum value along the given dimension.
1110 ///
1111 /// # Arguments
1112 ///
1113 /// * `dim` - The dimension or axis along which to aggregate the elements;
1114 /// supports negative indexing.
1115 ///
1116 /// # Returns
1117 ///
1118 /// The returned tensor will have the same rank,
1119 /// but the aggregated dimension will have size 1.
1120 ///
1121 /// # Example
1122 ///
1123 /// ```rust
1124 /// use burn_tensor::{Tensor, Shape};
1125 ///
1126 /// fn example() {
1127 /// let device = Default::default();
1128 /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
1129 /// let tensor = tensor.max_dim(0);
1130 /// println!("{tensor}");
1131 /// // [[5.0, 9.0, 6.0]]
1132 /// }
1133 /// ```
1134 pub fn max_dim<I: AsIndex>(self, dim: I) -> Self {
1135 let dim = unwrap_dim_index(dim.try_dim_index(D), "Max Dim");
1136 Tensor::new(K::max_dim(self.primitive, dim))
1137 }
1138
1139 /// Find the maximum value along the given dimensions.
1140 ///
1141 /// # Arguments
1142 ///
1143 /// * `dims` - The dimensions or axis along which to aggregate the elements;
1144 /// supports negative indexing.
1145 ///
1146 /// # Returns
1147 ///
1148 /// The returned tensor will have the same rank,
1149 /// but the aggregated dimensions will have size 1.
1150 ///
1151 /// # Example
1152 ///
1153 /// ```rust
1154 /// use burn_tensor::{Tensor, Shape};
1155 ///
1156 /// fn example() {
1157 /// let device = Default::default();
1158 /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
1159 /// let tensor = tensor.max_dims(&[0, 1]);
1160 /// println!("{tensor}");
1161 /// // [[9.0]]
1162 /// }
1163 /// ```
1164 pub fn max_dims<I: AsIndex>(self, dims: &[I]) -> Self {
1165 dims.iter().fold(self, |tensor, &dim| tensor.max_dim(dim))
1166 }
1167}