burn_tensor/tensor/api/numeric.rs
1use burn_backend::Scalar;
2
3use crate::alloc::borrow::ToOwned;
4use crate::check::unwrap_dim_index;
5use crate::kind::Numeric;
6
7use crate::{
8 AsIndex, Bool, Distribution, ElementConversion, Int, Shape, Tensor, check, check::TensorCheck,
9};
10use crate::{Device, IndexingUpdateOp, TensorCreationOptions};
11
12impl<const D: usize, K> Tensor<D, K>
13where
14 K: Numeric,
15{
16 /// Applies element wise addition operation.
17 ///
18 /// `y = x2 + x1`
19 ///
20 /// # Arguments
21 ///
22 /// * `other` - The tensor to add.
23 ///
24 /// # Example
25 ///
26 /// ```rust
27 /// use burn_tensor::{Tensor, Shape};
28 ///
29 /// fn example() {
30 /// let device = Default::default();
31 /// let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
32 /// let tensor2 = Tensor::<2>::from_data([[2.0, 3.0, 4.0], [1.0, 2.0, 3.0]], &device);
33 /// let tensor = tensor1 + tensor2;
34 /// println!("{tensor}");
35 /// // [[3.0, 1.0, 7.0], [6.0, 11.0, 9.0]]
36 /// }
37 /// ```
38 #[allow(clippy::should_implement_trait)]
39 pub fn add(self, other: Self) -> Self {
40 check!(TensorCheck::binary_ops_ew("Add", &self, &other));
41 Self::new(K::add(self.primitive, other.primitive))
42 }
43
44 /// Applies element wise addition operation with a scalar.
45 ///
46 /// `y = x + s`
47 ///
48 /// # Arguments
49 ///
50 /// * `other` - The scalar to add, element wise.
51 ///
52 /// # Example
53 ///
54 /// ```rust
55 /// use burn_tensor::{Tensor, Shape};
56 ///
57 /// fn example() {
58 /// let device = Default::default();
59 /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
60 /// let scalar = 2.0;
61 /// let tensor = tensor + scalar;
62 /// println!("{tensor}");
63 /// // [[3.0, 0.0, 5.0], [7.0, 11.0, 8.0]]
64 /// }
65 /// ```
66 pub fn add_scalar<E: ElementConversion>(self, other: E) -> Self {
67 let other = Scalar::new(other, &self.dtype());
68 Self::new(K::add_scalar(self.primitive, other))
69 }
70
71 /// Applies element wise subtraction operation.
72 ///
73 /// `y = x2 - x1`
74 ///
75 /// # Arguments
76 ///
77 /// * `other` - The tensor to subtract.
78 ///
79 /// # Example
80 ///
81 /// ```rust
82 /// use burn_tensor::{Tensor, Shape};
83 ///
84 /// fn example() {
85 /// let device = Default::default();
86 /// let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
87 /// let tensor2 = Tensor::<2>::from_data([[2.0, 3.0, 4.0], [1.0, 2.0, 3.0]], &device);
88 /// let tensor = tensor1 - tensor2;
89 /// println!("{tensor}");
90 /// // [[-1.0, -5.0, -1.0], [4.0, 7.0, 3.0]]
91 /// }
92 /// ```
93 #[allow(clippy::should_implement_trait)]
94 pub fn sub(self, other: Self) -> Self {
95 check!(TensorCheck::binary_ops_ew("Sub", &self, &other));
96 Self::new(K::sub(self.primitive, other.primitive))
97 }
98
99 /// Applies element wise subtraction operation with a scalar.
100 ///
101 /// `y = x - s`
102 ///
103 /// # Arguments
104 ///
105 /// * `other` - The scalar to subtract, element wise.
106 ///
107 /// # Example
108 ///
109 /// ```rust
110 /// use burn_tensor::{Tensor, Shape};
111 ///
112 /// fn example() {
113 /// let device = Default::default();
114 /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
115 /// let scalar = 2.0;
116 /// let tensor = tensor - scalar;
117 /// println!("{tensor}");
118 /// // [[-1.0, -4.0, 1.0], [3.0, 7.0, 4.0]]
119 /// }
120 /// ```
121 pub fn sub_scalar<E: ElementConversion>(self, other: E) -> Self {
122 let other = Scalar::new(other, &self.dtype());
123 Self::new(K::sub_scalar(self.primitive, other))
124 }
125
126 /// Applies element wise division operation.
127 ///
128 /// `y = x2 / x1`
129 ///
130 /// # Arguments
131 ///
132 /// * `other` - The tensor to divide.
133 ///
134 /// # Example
135 ///
136 /// ```rust
137 /// use burn_tensor::{Tensor, Shape};
138 ///
139 /// fn example() {
140 /// let device = Default::default();
141 /// let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
142 /// let tensor2 = Tensor::<2>::from_data([[2.0, 3.0, 4.0], [1.0, 2.0, 3.0]], &device);
143 /// let tensor = tensor1 / tensor2;
144 /// println!("{tensor}");
145 /// // [[0.5, -0.6666667, 0.75], [5.0, 4.5, 2.0]]
146 /// }
147 /// ```
148 #[allow(clippy::should_implement_trait)]
149 pub fn div(self, other: Self) -> Self {
150 check!(TensorCheck::binary_ops_ew("Div", &self, &other));
151 Self::new(K::div(self.primitive, other.primitive))
152 }
153
154 /// Applies element wise division operation with a scalar.
155 ///
156 /// `y = x / s`
157 ///
158 /// # Arguments
159 ///
160 /// * `other` - The scalar to divide, element wise.
161 ///
162 /// # Example
163 ///
164 /// ```rust
165 /// use burn_tensor::{Tensor, Shape};
166 ///
167 /// fn example() {
168 /// let device = Default::default();
169 /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
170 /// let scalar = 2.0;
171 /// let tensor = tensor / scalar;
172 /// println!("{tensor}");
173 /// // [[0.5, -1.0, 1.5], [2.5, 4.5, 3.0]]
174 /// }
175 /// ```
176 pub fn div_scalar<E: ElementConversion>(self, other: E) -> Self {
177 let other = Scalar::new(other, &self.dtype());
178 Self::new(K::div_scalar(self.primitive, other))
179 }
180
181 /// Applies element wise the remainder operation with a scalar.
182 ///
183 /// `y = x2 % x1`
184 pub fn remainder(self, other: Self) -> Self {
185 Self::new(K::remainder(self.primitive, other.primitive))
186 }
187
188 /// Applies element wise the remainder operation with a scalar.
189 ///
190 /// `y = x % s`
191 ///
192 /// # Arguments
193 ///
194 /// * `other` - The scalar to divide, element wise.
195 ///
196 /// # Example
197 ///
198 /// ```rust
199 /// use burn_tensor::{Tensor, Shape};
200 ///
201 /// fn example() {
202 /// let device = Default::default();
203 /// let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
204 /// let scalar = 2.0;
205 /// let tensor = tensor1 % scalar;
206 /// println!("{tensor}");
207 /// // [[1.0, 0.0, 1.0], [1.0, 1.0, 0.0]]
208 /// }
209 /// ```
210 pub fn remainder_scalar<E: ElementConversion>(self, other: E) -> Self {
211 let other = Scalar::new(other, &self.dtype());
212 Self::new(K::remainder_scalar(self.primitive, other))
213 }
214
215 /// Applies element wise multiplication operation.
216 ///
217 /// `y = x2 * x1`
218 ///
219 /// # Arguments
220 ///
221 /// * `other` - The tensor to multiply.
222 ///
223 /// # Example
224 ///
225 /// ```rust
226 /// use burn_tensor::{Tensor, Shape};
227 ///
228 /// fn example() {
229 /// let device = Default::default();
230 /// let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
231 /// let tensor2 = Tensor::<2>::from_data([[2.0, 3.0, 4.0], [1.0, 2.0, 3.0]], &device);
232 /// let tensor = tensor1 * tensor2;
233 /// println!("{tensor}");
234 /// // [[2.0, -6.0, 12.0], [5.0, 18.0, 18.0]]
235 /// }
236 /// ```
237 #[allow(clippy::should_implement_trait)]
238 pub fn mul(self, other: Self) -> Self {
239 check!(TensorCheck::binary_ops_ew("Mul", &self, &other));
240 Self::new(K::mul(self.primitive, other.primitive))
241 }
242
243 /// Applies element wise multiplication operation with a scalar.
244 ///
245 /// `y = x * s`
246 ///
247 /// # Arguments
248 ///
249 /// * `other` - The scalar to multiply, element wise.
250 ///
251 /// # Example
252 ///
253 /// ```rust
254 /// use burn_tensor::{Tensor, Shape};
255 ///
256 /// fn example() {
257 /// let device = Default::default();
258 /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
259 /// let scalar = 2.0;
260 /// let tensor = tensor * scalar;
261 /// println!("{tensor}");
262 /// // [[2.0, -4.0, 6.0], [10.0, 18.0, 12.0]]
263 /// }
264 /// ```
265 pub fn mul_scalar<E: ElementConversion>(self, other: E) -> Self {
266 let other = Scalar::new(other, &self.dtype());
267 Self::new(K::mul_scalar(self.primitive, other))
268 }
269
270 /// Switch sign of each element in the tensor.
271 ///
272 /// `y = -x`
273 ///
274 /// # Example
275 ///
276 /// ```rust
277 /// use burn_tensor::{Tensor, Shape};
278 ///
279 /// fn example() {
280 /// let device = Default::default();
281 /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
282 /// let tensor = -tensor;
283 /// println!("{tensor}");
284 /// // [[-1.0, 2.0, -3.0], [-5.0, -9.0, -6.0]]
285 /// }
286 /// ```
287 #[allow(clippy::should_implement_trait)]
288 pub fn neg(self) -> Self {
289 Self::new(K::neg(self.primitive))
290 }
291
292 /// Returns the signs of the elements of the input tensor.
293 ///
294 /// # Example
295 ///
296 /// ```rust
297 /// use burn_tensor::{Tensor, Shape};
298 ///
299 /// fn example() {
300 /// let device = Default::default();
301 /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
302 /// let tensor = tensor.sign();
303 /// println!("{tensor}");
304 /// // [[1.0, -1.0, 1.0], [1.0, 1.0, 1.0]]
305 /// }
306 /// ```
307 pub fn sign(self) -> Self {
308 Self::new(K::sign(self.primitive))
309 }
310
311 /// Aggregate all elements in the tensor with the mean operation.
312 ///
313 /// # Example
314 ///
315 /// ```rust
316 /// use burn_tensor::{Tensor, Shape};
317 ///
318 /// fn example() {
319 /// let device = Default::default();
320 /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
321 /// let tensor = tensor.mean();
322 /// println!("{tensor}");
323 /// // [3.6666667]
324 /// }
325 /// ```
326 pub fn mean(self) -> Tensor<1, K> {
327 Tensor::new(K::mean(self.primitive))
328 }
329
330 /// Aggregate all elements in the tensor with the sum operation.
331 ///
332 /// # Example
333 ///
334 /// ```rust
335 /// use burn_tensor::{Tensor, Shape};
336 ///
337 /// fn example() {
338 /// let device = Default::default();
339 /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
340 /// let tensor = tensor.sum();
341 /// println!("{tensor}");
342 /// // [22.0]
343 /// }
344 /// ```
345 pub fn sum(self) -> Tensor<1, K> {
346 Tensor::new(K::sum(self.primitive))
347 }
348
349 /// Aggregate all elements along the given *dimension* or *axis*
350 /// in the tensor with the mean operation.
351 ///
352 /// # Arguments
353 ///
354 /// * `dim` - The dimension or axis along which to aggregate the elements;
355 /// supports negative indexing.
356 ///
357 /// # Example
358 ///
359 /// ```rust
360 /// use burn_tensor::{Tensor, Shape};
361 ///
362 /// fn example() {
363 /// let device = Default::default();
364 /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
365 /// let tensor = tensor.clone().mean_dim(0);
366 /// println!("{tensor}");
367 /// // [[3.0, 3.5, 4.5]]
368 /// let tensor = tensor.clone().mean_dim(1);
369 /// println!("{tensor}");
370 /// // [[0.6666667], [6.6666665]]
371 /// }
372 /// ```
373 pub fn mean_dim<I: AsIndex>(self, dim: I) -> Self {
374 let dim = unwrap_dim_index(dim.try_dim_index(D), "Mean Dim");
375 Self::new(K::mean_dim(self.primitive, dim))
376 }
377
378 /// Aggregate all elements along the given *axes*
379 /// in the tensor with the mean operation.
380 ///
381 /// # Arguments
382 ///
383 /// * `dims` - the dimensions to aggregate; supports negative indexing.
384 ///
385 /// # Returns
386 ///
387 /// The returned tensor will have the same rank,
388 /// but the aggregated dimensions will have size 1.
389 ///
390 /// # Example
391 ///
392 /// ```rust
393 /// use burn_tensor::{Tensor, Shape};
394 ///
395 /// fn example() {
396 /// let device = Default::default();
397 /// let tensor = Tensor::<2>::from_data([[2.0, 4.0], [6.0, -4.0]], &device);
398 /// let tensor = tensor.clone().mean_dims(&[0, 1]);
399 /// println!("{tensor}");
400 /// // [[2.0]]
401 /// }
402 /// ```
403 pub fn mean_dims<I: AsIndex>(self, dims: &[I]) -> Self {
404 dims.iter().fold(self, |tensor, &dim| tensor.mean_dim(dim))
405 }
406
407 /// Aggregate all elements along the given *dimension* or *axis*
408 /// in the tensor with the sum operation.
409 ///
410 /// # Arguments
411 ///
412 /// * `dim` - The dimension or axis along which to aggregate the elements;
413 /// supports negative indexing.
414 ///
415 /// # Example
416 ///
417 /// ```rust
418 /// use burn_tensor::{Tensor, Shape};
419 ///
420 /// fn example() {
421 /// let device = Default::default();
422 /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
423 /// let tensor = tensor.clone().sum_dim(0);
424 /// println!("{tensor}");
425 /// // [[6.0, 7.0, 9.0]]
426 /// let tensor = tensor.clone().sum_dim(1);
427 /// println!("{tensor}");
428 /// // [[2.0], [20.0]]
429 /// }
430 /// ```
431 pub fn sum_dim<I: AsIndex>(self, dim: I) -> Self {
432 let dim = unwrap_dim_index(dim.try_dim_index(D), "Sum Dim");
433 Self::new(K::sum_dim(self.primitive, dim))
434 }
435
436 /// Aggregate all elements along the given *axes*
437 /// in the tensor with the sum operation.
438 ///
439 /// # Arguments
440 ///
441 /// * `dims` - the dimensions to aggregate; supports negative indexing.
442 ///
443 /// # Returns
444 ///
445 /// The returned tensor will have the same rank,
446 /// but the aggregated dimensions will have size 1.
447 ///
448 /// # Example
449 ///
450 /// ```rust
451 /// use burn_tensor::{Tensor, Shape};
452 ///
453 /// fn example() {
454 /// let device = Default::default();
455 /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
456 /// let tensor = tensor.clone().sum_dims(&[0, 1]);
457 /// println!("{tensor}");
458 /// // [[27]]
459 /// }
460 /// ```
461 pub fn sum_dims<I: AsIndex>(self, dims: &[I]) -> Self {
462 dims.iter().fold(self, |tensor, &dim| tensor.sum_dim(dim))
463 }
464
465 /// Aggregate and squeeze along the given dimensions.
466 ///
467 /// This is equivalent to ``tensor.sum_dims(dims).squeeze_dims(dims)``
468 ///
469 /// # Arguments
470 ///
471 /// * `dims` - the dimensions to aggregate; supports negative indexing.
472 ///
473 /// # Returns
474 ///
475 /// The returned tensor will have the same rank,
476 /// but the aggregated dimensions will have size 1.
477 ///
478 /// # Example
479 ///
480 /// ```rust
481 /// use burn_tensor::{Tensor, Shape};
482 ///
483 /// fn example() {
484 /// let device = Default::default();
485 /// let tensor = Tensor::<3>::from_data([
486 /// [[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]],
487 /// [[9.0, 2.0, 5.0], [5.0, 7.0, 7.0]],
488 /// ], &device);
489 /// let tensor = tensor.clone().sum_dims_squeeze::<1, _>(&[0, 1]);
490 /// println!("{tensor}");
491 /// // [20.0, 16.0, 21.0]
492 /// }
493 /// ```
494 pub fn sum_dims_squeeze<const D2: usize, I: AsIndex>(self, dims: &[I]) -> Tensor<D2, K> {
495 self.sum_dims(dims).squeeze_dims::<D2>(dims)
496 }
497
498 /// Aggregate all elements in the tensor with the product operation.
499 ///
500 /// # Example
501 ///
502 /// ```rust
503 /// use burn_tensor::{Tensor, Shape};
504 ///
505 /// fn example() {
506 /// let device = Default::default();
507 /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
508 /// let tensor = tensor.prod();
509 /// println!("{tensor}");
510 /// // [-1620.0]
511 /// }
512 /// ```
513 pub fn prod(self) -> Tensor<1, K> {
514 Tensor::new(K::prod(self.primitive))
515 }
516
517 /// Aggregate all elements along the given *dimension* or *axis*
518 /// in the tensor with the product operation.
519 ///
520 /// # Arguments
521 ///
522 /// * `dim` - The dimension or axis along which to aggregate the elements,
523 /// supports negative indexing.
524 ///
525 /// # Returns
526 ///
527 /// The returned tensor will have the same rank,
528 /// but the aggregated dimension will have size 1.
529 ///
530 /// # Example
531 ///
532 /// ```rust
533 /// use burn_tensor::{Tensor, Shape};
534 ///
535 /// fn example() {
536 /// let device = Default::default();
537 /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
538 /// let tensor = tensor.clone().prod_dim(0);
539 /// println!("{tensor}");
540 /// // [[5.0, -18.0, 18.0]]
541 /// let tensor = tensor.clone().prod_dim(1);
542 /// println!("{tensor}");
543 /// // [[-6.0], [270.0]]
544 /// }
545 /// ```
546 pub fn prod_dim<I: AsIndex>(self, dim: I) -> Self {
547 let dim = unwrap_dim_index(dim.try_dim_index(D), "Prod Dim");
548 Self::new(K::prod_dim(self.primitive, dim))
549 }
550
551 /// Aggregate all elements along the given *axes*
552 /// in the tensor with the prod operation.
553 ///
554 /// # Arguments
555 ///
556 /// * `dims` - the dimensions to aggregate, supports negative indexing.
557 ///
558 /// # Returns
559 ///
560 /// The returned tensor will have the same rank,
561 /// but the aggregated dimensions will have size 1.
562 ///
563 /// # Example
564 ///
565 /// ```rust
566 /// use burn_tensor::{Tensor, Shape};
567 ///
568 /// fn example() {
569 /// let device = Default::default();
570 /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
571 /// let tensor = tensor.clone().sum_dims(&[0, 1]);
572 /// println!("{tensor}");
573 /// // [[-1620.0]]
574 /// }
575 /// ```
576 pub fn prod_dims<I: AsIndex>(self, dims: &[I]) -> Self {
577 dims.iter().fold(self, |tensor, &dim| tensor.prod_dim(dim))
578 }
579
580 /// Computes the cumulative sum of elements along the given *dimension* or *axis*.
581 ///
582 /// # Arguments
583 ///
584 /// * `dim` - The dimension or axis along which to compute the cumulative sum.
585 /// Negative dimensions are supported and count from the end.
586 ///
587 /// # Example
588 ///
589 /// ```rust
590 /// use burn_tensor::{Tensor, Shape};
591 ///
592 /// fn example() {
593 /// let device = Default::default();
594 /// let tensor = Tensor::<2>::from_data([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], &device);
595 /// let result = tensor.clone().cumsum(0);
596 /// println!("{result}");
597 /// // [[1.0, 2.0, 3.0], [5.0, 7.0, 9.0]]
598 /// let result = tensor.cumsum(1);
599 /// println!("{result}");
600 /// // [[1.0, 3.0, 6.0], [4.0, 9.0, 15.0]]
601 /// }
602 /// ```
603 pub fn cumsum<I: AsIndex>(self, dim: I) -> Self {
604 let dim = unwrap_dim_index(dim.try_dim_index(D), "Cumsum");
605 Self::new(K::cumsum(self.primitive, dim))
606 }
607
608 /// Computes the cumulative product of elements along the given *dimension* or *axis*.
609 ///
610 /// # Arguments
611 ///
612 /// * `dim` - The dimension or axis along which to compute the cumulative product.
613 /// Negative dimensions are supported and count from the end.
614 ///
615 /// # Example
616 ///
617 /// ```rust
618 /// use burn_tensor::{Tensor, Shape};
619 ///
620 /// fn example() {
621 /// let device = Default::default();
622 /// let tensor = Tensor::<2>::from_data([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], &device);
623 /// let result = tensor.clone().cumprod(0);
624 /// println!("{result}");
625 /// // [[1.0, 2.0, 3.0], [4.0, 10.0, 18.0]]
626 /// let result = tensor.cumprod(1);
627 /// println!("{result}");
628 /// // [[1.0, 2.0, 6.0], [4.0, 20.0, 120.0]]
629 /// }
630 /// ```
631 pub fn cumprod<I: AsIndex>(self, dim: I) -> Self {
632 let dim = unwrap_dim_index(dim.try_dim_index(D), "Cumprod");
633 Self::new(K::cumprod(self.primitive, dim))
634 }
635
636 /// Apply element wise absolute value operation.
637 ///
638 /// # Example
639 ///
640 /// ```rust
641 /// use burn_tensor::{Int, Tensor};
642 ///
643 /// fn example() {
644 /// let device = Default::default();
645 /// let tensor = Tensor::<2, Int>::from_ints([[1, -2, 3], [4, -5, 6], [7, -8, 9]], &device);
646 /// let tensor = tensor.abs();
647 /// println!("{tensor}");
648 /// // [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
649 /// }
650 /// ```
651 ///
652 /// # Notes
653 ///
654 /// For signed integer dtypes, this operation uses two's-complement wraparound semantics, similar to
655 /// `x.wrapping_abs()`. For example, `abs(i64::MIN) == i64::MIN`.
656 pub fn abs(self) -> Self {
657 Self::new(K::abs(self.primitive))
658 }
659
660 /// Returns the upper triangular part of a matrix (2-D tensor) or batch of matrices input,
661 /// the other elements of the result tensor out are set to 0.
662 ///
663 /// See also [`triu_mask`](Tensor::triu_mask).
664 ///
665 /// # Arguments
666 ///
667 /// * `diagonal` - The offset from the diagonal, where 0 means the diagonal, and positive values shift
668 /// towards the upper triangle.
669 ///
670 /// # Example
671 /// ```rust
672 /// use burn_tensor::{Int, Tensor};
673 ///
674 /// fn example() {
675 /// let device = Default::default();
676 /// let tensor = Tensor::<2, Int>::from_ints(
677 /// [
678 /// [1, 2, 3],
679 /// [4, 5, 6],
680 /// [7, 8, 9]
681 /// ],
682 /// &device
683 /// );
684 /// let tensor = tensor.triu(1);
685 /// println!("{tensor}");
686 /// // [
687 /// // [0, 2, 3],
688 /// // [0, 0, 6],
689 /// // [0, 0, 0]
690 /// // ]
691 /// }
692 /// ```
693 pub fn triu(self, diagonal: i64) -> Self {
694 check!(TensorCheck::tri::<{ D }>());
695
696 // last two dimensions
697 let shape = &self.shape()[D - 2..].to_owned();
698
699 let mask = Tensor::<2, Bool>::triu_mask(shape, diagonal, &self.device()).unsqueeze();
700 self.mask_fill(mask, 0)
701 }
702
703 /// Returns the lower triangular part of a matrix (2-D tensor) or batch of matrices input,
704 /// the other elements of the result tensor out are set to 0.
705 ///
706 /// See also [`tril_mask`](Tensor::tril_mask).
707 ///
708 /// # Arguments
709 ///
710 /// * `diagonal` - The offset from the diagonal, where 0 means the diagonal, and positive values shift
711 /// towards the upper triangle.
712 ///
713 /// # Example
714 /// ```rust
715 /// use burn_tensor::{Int, Tensor};
716 ///
717 /// fn example() {
718 /// let device = Default::default();
719 /// let tensor = Tensor::<2, Int>::from_ints(
720 /// [
721 /// [1, 2, 3],
722 /// [4, 5, 6],
723 /// [7, 8, 9]
724 /// ],
725 /// &device
726 /// );
727 ///
728 /// let tensor = tensor.tril(-1);
729 /// println!("{tensor}");
730 /// // [
731 /// // [0, 0, 0],
732 /// // [4, 0, 0],
733 /// // [7, 8, 0]
734 /// // ]
735 /// }
736 /// ```
737 pub fn tril(self, diagonal: i64) -> Self {
738 check!(TensorCheck::tri::<{ D }>());
739
740 // last two dimensions
741 let shape = &self.shape()[D - 2..].to_owned();
742 let mask = Tensor::<2, Bool>::tril_mask(shape, diagonal, &self.device()).unsqueeze();
743
744 self.mask_fill(mask, 0)
745 }
746
747 /// Applies element wise power operation with a integer Tensor
748 ///
749 /// # Arguments
750 ///
751 /// * `other` - The tensor to apply the power operation with.
752 ///
753 /// # Example
754 ///
755 /// ```rust
756 /// use burn_tensor::{Tensor, Shape, Int};
757 ///
758 /// fn example() {
759 /// let device = Default::default();
760 /// let tensor1 = Tensor::<2, Int>::from_ints([[1, -2, 3], [5, 9, 6]], &device);
761 /// let tensor2 = Tensor::<2, Int>::from_ints([[2, 3, 4], [1, 2, 3]], &device);
762 /// let tensor = tensor1.powi(tensor2);
763 /// println!("{tensor}");
764 /// // [[1, -8, 81], [5, 81, 216]]
765 /// }
766 /// ```
767 pub fn powi(self, other: Self) -> Self {
768 Self::new(K::powi(self.primitive, other.primitive))
769 }
770
771 /// Applies element wise power operation with a integer scalar
772 ///
773 /// # Arguments
774 ///
775 /// * `other` - The scalar to apply the power operation with.
776 ///
777 /// # Example
778 ///
779 /// ```rust
780 /// use burn_tensor::{Tensor, Shape, Int};
781 ///
782 /// fn example() {
783 /// let device = Default::default();
784 /// let tensor = Tensor::<2, Int>::from_ints([[1, -2, 3], [5, 9, 6]], &device);
785 /// let tensor = tensor.powi_scalar(2);
786 /// println!("{tensor}");
787 ///
788 /// // [[1, 4, 9], [25, 81, 36]]
789 /// let tensor = Tensor::<2>::from_data([[1.5, -2., 3.], [5., 9., 6.]], &device);
790 /// let tensor = tensor.powi_scalar(2);
791 /// println!("{tensor}");
792 /// // [[2.25, 4., 9.], [25., 81., 36.]]
793 /// }
794 /// ```
795 pub fn powi_scalar<E: ElementConversion>(self, other: E) -> Self {
796 let other = Scalar::new(other, &self.dtype());
797 Self::new(K::powi_scalar(self.primitive, other))
798 }
799
800 /// Converts the tensor to a boolean tensor by checking if the elements are non-zero.
801 ///
802 /// # Returns
803 ///
804 /// A boolean tensor with the same shape as the input tensor.
805 ///
806 /// # Example
807 ///
808 /// ```rust
809 /// use burn_tensor::{Tensor, Shape};
810 ///
811 /// fn example() {
812 /// let device = Default::default();
813 /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [0.0, 9.0, 6.0]], &device);
814 /// let tensor = tensor.bool();
815 /// println!("{tensor}");
816 /// // [
817 /// // [true, true, true],
818 /// // [false, true, true]
819 /// // ]
820 /// }
821 pub fn bool(self) -> Tensor<D, Bool> {
822 self.not_equal_scalar(0)
823 }
824
825 /// Create a random tensor of the given shape on the given device where each element is
826 /// sampled from the given distribution.
827 ///
828 /// See also [`random_like`](Tensor::random_like).
829 ///
830 /// # Arguments
831 ///
832 /// * `shape` - The shape of the tensor.
833 /// * `distribution` - The distribution to sample from.
834 /// * `device` - The device to create the tensor on.
835 ///
836 /// # Returns
837 ///
838 /// A new tensor with the given shape and elements sampled from the given distribution.
839 ///
840 /// # Example
841 ///
842 /// ```rust
843 /// use burn_tensor::{Tensor, Shape, Distribution};
844 ///
845 /// fn example() {
846 /// let device = Default::default();
847 /// let distribution = Distribution::Uniform(0.0, 1.0); // Any random value between 0.0 and 1.0
848 /// let tensor = Tensor::<2>::random(Shape::new([2, 3]), distribution, &device);
849 /// println!("{tensor}");
850 /// // [
851 /// // [0.08347523, 0.70498955, 0.60332155],
852 /// // [0.08173251, 0.18028641, 0.97942924]
853 /// // ]
854 /// }
855 /// ```
856 pub fn random<S: Into<Shape>>(
857 shape: S,
858 distribution: Distribution,
859 options: impl Into<TensorCreationOptions>,
860 ) -> Self {
861 // Use the given dtype when provided, otherwise default device dtype
862 let opt = options.into();
863 let dtype = opt.resolve_dtype::<K>();
864 Self::new(K::random(shape.into(), distribution, &opt.device, dtype))
865 }
866
867 /// Applies the matrix multiplication operation.
868 ///
869 /// ```math
870 /// C = AB
871 /// ```
872 ///
873 /// Shapes of the form `[..., B, 1, K] @ [..., 1, K, N]` are reinterpreted as
874 /// `[..., 1, B, K] @ [..., 1, K, N]`, turning a batched vec-mat into a general
875 /// matmul, which is often faster.
876 pub fn matmul(self, other: Self) -> Self {
877 check!(TensorCheck::matmul(&self, &other));
878
879 if D >= 3 {
880 let batch_index = D - 3;
881 let vector_index = D - 2;
882 let lhs_dims = &self.shape()[batch_index..D];
883 let rhs_dims = &other.shape()[batch_index..D];
884
885 if let ([_, 1, k1], [1, k2, _]) = (lhs_dims, rhs_dims)
886 && k1 == k2
887 {
888 return Tensor::new(K::matmul(
889 self.swap_dims(batch_index, vector_index).primitive,
890 other.primitive,
891 ))
892 .swap_dims(batch_index, vector_index);
893 }
894 }
895
896 Tensor::new(K::matmul(self.primitive, other.primitive))
897 }
898}
899
900impl<K> Tensor<1, K>
901where
902 K: Numeric,
903{
904 /// Calculates the dot product with another tensor.
905 ///
906 /// `y = x2.dot(x1)`
907 ///
908 /// # Arguments
909 ///
910 /// * `other` - The tensor to compute dot product with.
911 ///
912 /// # Notes
913 ///
914 /// Both tensors must have the same number of elements.
915 ///
916 /// # Example
917 ///
918 /// ```rust
919 /// use burn_tensor::{Tensor, Shape};
920 ///
921 /// fn example() {
922 /// let device = Default::default();
923 /// let tensor1 = Tensor::<1>::from_data([1.0, 2.0], &device);
924 /// let tensor2 = Tensor::<1>::from_data([-2.0, 3.0], &device);
925 /// let tensor = tensor1.dot(tensor2);
926 /// println!("{tensor}");
927 /// // [4]
928 /// }
929 /// ```
930 pub fn dot(self, other: Self) -> Self {
931 self.mul(other).sum()
932 }
933}
934
935impl<K> Tensor<2, K>
936where
937 K: Numeric,
938{
939 /// Creates a new 2D tensor with ones on the diagonal and zeros elsewhere.
940 ///
941 /// # Arguments
942 ///
943 /// * `size` - The size of the square matrix.
944 pub fn eye(size: usize, device: &Device) -> Self {
945 let indices = Tensor::<1, Int>::arange(0..size as i64, device).unsqueeze::<2>();
946 let ones = Self::ones([1, size], device);
947 let zeros = Self::zeros([size, size], device);
948
949 zeros.scatter(0, indices, ones, IndexingUpdateOp::Add)
950 }
951}
952
953// Tensor + tensor
954impl<const D: usize, K: Numeric> core::ops::Add<Self> for Tensor<D, K> {
955 type Output = Self;
956
957 fn add(self, rhs: Self) -> Self::Output {
958 Self::add(self, rhs)
959 }
960}
961
962// Tensor + scalar
963impl<E: ElementConversion, const D: usize, K: Numeric> core::ops::Add<E> for Tensor<D, K> {
964 type Output = Self;
965
966 fn add(self, other: E) -> Self::Output {
967 Tensor::add_scalar(self, other)
968 }
969}
970
971// Scalar + tensor
972macro_rules! impl_tensor_scalar_add {
973 ($($t:ty),*) => {
974 $(
975 impl<const D: usize, K: Numeric> core::ops::Add<Tensor<D, K>> for $t
976 {
977 type Output = Tensor<D, K>;
978
979 fn add(self, tensor: Tensor<D, K>) -> Self::Output {
980 Tensor::add_scalar(tensor, self)
981 }
982 }
983 )*
984 }
985}
986impl_tensor_scalar_add!(f32, f64, i32, i64, u32, u64);
987
988// Tensor - tensor
989impl<const D: usize, K: Numeric> core::ops::Sub<Self> for Tensor<D, K> {
990 type Output = Self;
991
992 fn sub(self, rhs: Self) -> Self::Output {
993 Tensor::sub(self, rhs)
994 }
995}
996
997// Tensor - scalar
998impl<E: ElementConversion, const D: usize, K: Numeric> core::ops::Sub<E> for Tensor<D, K> {
999 type Output = Self;
1000
1001 fn sub(self, other: E) -> Self::Output {
1002 Tensor::sub_scalar(self, other)
1003 }
1004}
1005
1006// Scalar - tensor
1007macro_rules! impl_tensor_scalar_sub {
1008 ($($t:ty),*) => {
1009 $(
1010 impl<const D: usize, K: Numeric> core::ops::Sub<Tensor<D, K>> for $t
1011 {
1012 type Output = Tensor<D, K>;
1013
1014 fn sub(self, tensor: Tensor<D, K>) -> Self::Output {
1015 Tensor::add_scalar(Tensor::neg(tensor), self)
1016 }
1017 }
1018 )*
1019 }
1020}
1021impl_tensor_scalar_sub!(f32, f64, i32, i64, u32, u64);
1022
1023// Tensor / tensor
1024impl<const D: usize, K: Numeric> core::ops::Div<Self> for Tensor<D, K> {
1025 type Output = Self;
1026
1027 fn div(self, rhs: Self) -> Self::Output {
1028 Tensor::div(self, rhs)
1029 }
1030}
1031
1032// Tensor / scalar
1033impl<E: ElementConversion, const D: usize, K: Numeric> core::ops::Div<E> for Tensor<D, K> {
1034 type Output = Self;
1035
1036 fn div(self, other: E) -> Self::Output {
1037 Tensor::div_scalar(self, other)
1038 }
1039}
1040
1041// Scalar / tensor (float only)
1042macro_rules! impl_tensor_scalar_div {
1043 ($($t:ty),*) => {
1044 $(
1045 impl<const D: usize> core::ops::Div<Tensor<D>> for $t
1046 {
1047 type Output = Tensor<D>;
1048
1049 fn div(self, tensor: Tensor<D>) -> Self::Output {
1050 tensor.recip().mul_scalar(self)
1051 }
1052 }
1053 )*
1054 }
1055}
1056
1057impl_tensor_scalar_div!(f32, f64);
1058
1059// Tensor % tensor.
1060impl<const D: usize, K: Numeric> core::ops::Rem<Self> for Tensor<D, K> {
1061 type Output = Self;
1062
1063 fn rem(self, rhs: Self) -> Self::Output {
1064 Tensor::remainder(self, rhs)
1065 }
1066}
1067
1068// Tensor % scalar.
1069impl<E: ElementConversion, const D: usize, K: Numeric> core::ops::Rem<E> for Tensor<D, K> {
1070 type Output = Self;
1071
1072 fn rem(self, other: E) -> Self::Output {
1073 Tensor::remainder_scalar(self, other)
1074 }
1075}
1076
1077// Tensor * tensor.
1078impl<const D: usize, K: Numeric> core::ops::Mul<Self> for Tensor<D, K> {
1079 type Output = Self;
1080
1081 fn mul(self, rhs: Self) -> Self::Output {
1082 Tensor::mul(self, rhs)
1083 }
1084}
1085
1086// Tensor * scalar.
1087impl<E: ElementConversion, const D: usize, K: Numeric> core::ops::Mul<E> for Tensor<D, K> {
1088 type Output = Self;
1089
1090 fn mul(self, other: E) -> Self::Output {
1091 Tensor::mul_scalar(self, other)
1092 }
1093}
1094
1095macro_rules! impl_tensor_scalar_mul {
1096 ($($t:ty),*) => {
1097 $(
1098 impl<const D: usize, K: Numeric> core::ops::Mul<Tensor<D, K>> for $t
1099 {
1100 type Output = Tensor<D, K>;
1101
1102 fn mul(self, other: Tensor<D, K>) -> Self::Output {
1103 Tensor::mul_scalar(other, self)
1104 }
1105 }
1106 )*
1107 }
1108}
1109
1110impl_tensor_scalar_mul!(f32, f64, i32, i64, u32, u64);
1111
1112impl<const D: usize, K> core::ops::Neg for Tensor<D, K>
1113where
1114 K: Numeric,
1115{
1116 type Output = Self;
1117
1118 fn neg(self) -> Self::Output {
1119 Tensor::neg(self)
1120 }
1121}