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