burn_tensor/tensor/api/
fmod.rs

1use crate::{Float, Tensor};
2
3impl<const D: usize> Tensor<D, Float> {
4    /// Computes the floating-point remainder of dividing `self` by `other`.
5    ///
6    /// The result has the same sign as `self` and magnitude less than `other`.
7    /// This is equivalent to the IEEE 754 remainder operation.
8    ///
9    /// # Special Cases (IEEE 754 compliant)
10    ///
11    /// - If `self` is ±∞ and `other` is not NaN, NaN is returned
12    /// - If `other` is ±0 and `self` is not NaN, NaN is returned
13    /// - If `other` is ±∞ and `self` is finite, `self` is returned
14    /// - If either argument is NaN, NaN is returned
15    ///
16    /// # Arguments
17    ///
18    /// * `other` - The divisor tensor. Must have the same shape as `self`.
19    ///
20    /// # Returns
21    ///
22    /// A tensor with the same shape where each element is the floating-point remainder.
23    ///
24    /// # Example
25    ///
26    /// ```rust
27    /// use burn_tensor::Tensor;
28    ///
29    /// let device = Default::default();
30    /// let dividend = Tensor::<1>::from_data([5.3, -5.3, 5.3, -5.3], &device);
31    /// let divisor = Tensor::<1>::from_data([2.0, 2.0, -2.0, -2.0], &device);
32    /// let result = dividend.fmod(divisor);
33    ///
34    /// // Result: [1.3, -1.3, 1.3, -1.3]
35    /// ```
36    pub fn fmod(self, other: Self) -> Self {
37        // Normal case: fmod(x, y) = x - y * trunc(x / y)
38        let quotient = self.clone().div(other.clone());
39        let truncated = quotient.trunc();
40        let product = other.clone() * truncated.clone();
41
42        // When divisor is infinity and dividend is finite:
43        // - quotient is 0, truncated is 0
44        // - but 0 * infinity = NaN, which is wrong
45        // We need to handle this case by replacing NaN with 0 when appropriate
46
47        // Check if the product is NaN due to 0 * inf
48        let is_zero_times_inf = truncated.equal_scalar(0.0).bool_and(other.is_inf());
49        let zero_tensor = self.clone().mul_scalar(0.0);
50        let corrected_product = product.mask_where(is_zero_times_inf, zero_tensor);
51
52        self - corrected_product
53    }
54
55    /// Computes the floating-point remainder of dividing `self` by a scalar.
56    ///
57    /// The result has the same sign as `self` and magnitude less than the scalar.
58    ///
59    /// # Special Cases (IEEE 754 compliant)
60    ///
61    /// - If `self` is ±∞ and scalar is not NaN, NaN is returned
62    /// - If scalar is ±0 and `self` is not NaN, NaN is returned
63    /// - If scalar is ±∞ and `self` is finite, `self` is returned
64    /// - If either argument is NaN, NaN is returned
65    ///
66    /// # Arguments
67    ///
68    /// * `scalar` - The scalar divisor.
69    ///
70    /// # Returns
71    ///
72    /// A tensor with the same shape where each element is the floating-point remainder.
73    ///
74    /// # Example
75    ///
76    /// ```rust
77    /// use burn_tensor::Tensor;
78    ///
79    /// let device = Default::default();
80    /// let tensor = Tensor::<1>::from_data([5.3, -5.3, 7.5, -7.5], &device);
81    /// let result = tensor.fmod_scalar(2.0);
82    ///
83    /// // Result: [1.3, -1.3, 1.5, -1.5]
84    /// ```
85    pub fn fmod_scalar(self, scalar: f32) -> Self {
86        // Normal case: fmod(x, y) = x - y * trunc(x / y)
87        let quotient = self.clone().div_scalar(scalar);
88        let truncated = quotient.trunc();
89        let product = truncated.mul_scalar(scalar);
90
91        // Handle the special case where scalar is infinity
92        // When scalar is ±∞ and self is finite, quotient is 0, truncated is 0
93        // but 0 * infinity = NaN, which is wrong - it should be 0
94        if scalar.is_infinite() {
95            // For finite values, fmod(x, ±∞) = x
96            // For infinite values, fmod(±∞, ±∞) = NaN (which is handled by arithmetic)
97            return self;
98        }
99
100        self - product
101    }
102}
burn_tensor/tensor/api/fmod.rs

burn_tensor/tensor/api/
fmod.rs
