burn_tensor/tensor/api/trunc.rs
1use burn_backend::ops::FloatTensorOps;
2use burn_dispatch::Dispatch;
3
4use crate::{Float, Tensor, ops::BridgeTensor};
5
6impl<const D: usize> Tensor<D, Float> {
7 /// Truncates the tensor element-wise, rounding toward zero.
8 ///
9 /// This function returns a new tensor with the same shape as the input tensor,
10 /// where each element is truncated toward zero. For positive values, this is
11 /// equivalent to floor, and for negative values, it's equivalent to ceil.
12 ///
13 /// # Special Cases (IEEE 754 compliant)
14 ///
15 /// - `trunc(±0)` returns ±0 (preserves sign of zero)
16 /// - `trunc(±∞)` returns ±∞
17 /// - `trunc(NaN)` returns NaN
18 ///
19 /// # Returns
20 ///
21 /// A tensor with the same shape where each element has been truncated toward zero.
22 ///
23 /// # Example
24 ///
25 /// ```rust
26 /// use burn_tensor::Tensor;
27 ///
28 /// fn example() {
29 /// let device = Default::default();
30 /// let tensor = Tensor::<1>::from_data([2.3, -1.7, 0.5, -0.5, 3.9], &device);
31 /// let truncated = tensor.trunc();
32 ///
33 /// // Result: [2.0, -1.0, 0.0, -0.0, 3.0]
34 /// }
35 /// ```
36 pub fn trunc(self) -> Self {
37 Self::new(trunc_impl(self.primitive))
38 }
39}
40
41fn trunc_impl(p: BridgeTensor) -> BridgeTensor {
42 BridgeTensor::float(Dispatch::float_trunc(p.into_float()))
43}