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