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

burn_tensor/tensor/api/
fmod.rs
