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