burn_tensor/tensor/api/

float.rs

use crate::AsIndex;
use crate::Cast;
use crate::Device;
use crate::Tensor;
use crate::cast::ToElement;
use crate::check;
use crate::check::TensorCheck;
use crate::check::unwrap_dim_index;
use crate::kind::FloatMath;
use crate::ops::{BridgeKind, BridgeTensor};
use crate::quantization::{QuantScheme, QuantizationParameters};
use crate::tensor::stats;
use crate::tensor::{Distribution, TensorData};
use crate::{Bool, Float, Int, TensorPrimitive};
#[cfg(feature = "std")]
use burn_backend::AutodiffBackend;
use burn_backend::ElementConversion;
use burn_backend::Scalar;
use burn_backend::TensorMetadata;
#[cfg(feature = "std")]
use burn_backend::distributed::DistributedParamId;
use burn_backend::ops::ActivationOps;
use burn_backend::ops::FloatTensorOps;
use burn_backend::ops::GridSampleOptions;
use burn_backend::ops::QTensorOps;
use burn_backend::quantization::QuantizationParametersPrimitive;
use burn_dispatch::Dispatch;
use core::f32;

/// Default RTOL value for `is_close` and `all_close`.
pub const DEFAULT_RTOL: f64 = 1e-5;

/// Default ATOL value for `is_close` and `all_close`.
pub const DEFAULT_ATOL: f64 = 1e-8;

impl<const D: usize> Tensor<D> {
    /// Applies the [error function](https://en.wikipedia.org/wiki/Error_function) element wise.
    ///
    #[cfg_attr(
        doc,
        doc = r#"
$y_i = \text{erf}\(x_i\)$

The error function is defined as:

$$\text{erf}\(x\) = \frac{2}{\sqrt{\pi}} \int_0^x e^{-t^2} dt$$
"#
    )]
    #[cfg_attr(not(doc), doc = "`y_i = erf(x_i)`")]
    pub fn erf(self) -> Self {
        Self::new(erf_impl(self.primitive))
    }

    /// Applies [hypotenuse operation](https://en.wikipedia.org/wiki/Hypotenuse) element wise.
    ///
    #[cfg_attr(doc, doc = r#"$y_i = \sqrt{x_i^2 + y_i^2}$"#)]
    #[cfg_attr(not(doc), doc = "`y_i = sqrt(x_i^2 + y_i^2)`")]
    pub fn hypot(self, other: Self) -> Self {
        Self::new(hypot_impl(self.primitive, other.primitive))
    }

    /// Applies [reciprocal operation](https://en.wikipedia.org/wiki/Multiplicative_inverse)
    /// (or multiplicative inverse) element wise.
    ///
    #[cfg_attr(doc, doc = r#"$y_i = \frac{1}{x_i}$"#)]
    #[cfg_attr(not(doc), doc = "`y_i = 1/x_i`")]
    pub fn recip(self) -> Self {
        Self::new(recip_impl(self.primitive))
    }

    /// Converts each of the elements of the input tensor from angles in degrees to radians.
    ///
    /// # Example
    /// ```ignore
    /// let tensor_in_radians = tensor.deg2rad();
    /// ```
    pub fn deg2rad(self) -> Self {
        self.mul_scalar(f32::consts::PI / 180.0)
    }

    /// Converts each of the elements of the input tensor from angles in radians to degrees.
    ///
    /// # Example
    /// ```ignore
    /// let tensor_in_degrees = tensor.rad2deg();
    /// ```
    pub fn rad2deg(self) -> Self {
        self.mul_scalar(180.0 / f32::consts::PI)
    }

    /// Applies element wise round operation.
    ///
    /// This function implements the [round half to even](https://en.wikipedia.org/wiki/Rounding#Rounding_half_to_even)
    /// strategy, with halfway cases rounded to the nearest even integer value.
    pub fn round(self) -> Self {
        Self::new(round_impl(self.primitive))
    }

    /// Applies element wise floor operation.
    pub fn floor(self) -> Self {
        Self::new(floor_impl(self.primitive))
    }

    /// Applies element wise ceil operation.
    pub fn ceil(self) -> Self {
        Self::new(ceil_impl(self.primitive))
    }

    /// Create a tensor from floats (f32) on a given device.
    ///
    /// # Example
    ///
    /// ```rust
    /// use burn_tensor::Tensor;
    ///
    /// let device = Default::default();
    /// let _ = Tensor::<1>::from_floats([1.0, 2.0], &device);
    /// let _ = Tensor::<2>::from_floats([[1.0, 2.0], [3.0, 4.0]], &device);
    /// ```
    pub fn from_floats<A: Into<TensorData>>(floats: A, device: &Device) -> Self {
        Self::from_data(floats.into().convert::<f32>(), device)
    }

    /// Returns a new tensor with the same shape and device as the current tensor and the data
    /// cast to Integer.
    ///
    /// # Example
    ///
    /// ```rust
    /// use burn_tensor::Tensor;
    ///
    /// let device = Default::default();
    /// let float_tensor = Tensor::<1>::from_floats([1.0, 2.0], &device);
    /// let int_tensor = float_tensor.int();
    /// ```
    pub fn int(self) -> Tensor<D, Int> {
        let device = self.device();
        Tensor::new(int_impl(self.primitive, device))
    }

    /// Returns a new tensor with the same shape, dtype, and device as the current tensor filled random
    /// values sampled from the given distribution.
    pub fn random_like(&self, distribution: Distribution) -> Self {
        Self::new(random_like_impl(&self.primitive, distribution))
    }

    /// Calculate the variance along the given dimension.
    ///
    /// Negative dimensions are supported and count from the end.
    pub fn var<I: AsIndex>(self, dim: I) -> Self {
        let dim = unwrap_dim_index(dim.try_dim_index(D), "Var");
        stats::var(self, dim)
    }

    /// Calculate the variance along the given dimension without applying the Bessel’s correction.
    ///
    /// Negative dimensions are supported and count from the end.
    pub fn var_bias<I: AsIndex>(self, dim: I) -> Self {
        let dim = unwrap_dim_index(dim.try_dim_index(D), "Var Bias");
        stats::var_bias(self, dim)
    }

    /// Calculate the variance along the given dimension and also returns the mean.
    ///
    /// Negative dimensions are supported and count from the end.
    pub fn var_mean<I: AsIndex>(self, dim: I) -> (Self, Self) {
        let dim = unwrap_dim_index(dim.try_dim_index(D), "Var Mean");
        let mean = self.clone().mean_dim(dim);
        let var = stats::var_with_mean(self, mean.clone(), dim);
        (var, mean)
    }

    /// Calculate the variance along the given dimension without applying the Bessel’s correction and also returns the mean.
    ///
    /// Negative dimensions are supported and count from the end.
    pub fn var_mean_bias<I: AsIndex>(self, dim: I) -> (Self, Self) {
        let dim = unwrap_dim_index(dim.try_dim_index(D), "Var Mean Bias");
        let mean = self.clone().mean_dim(dim);
        let var = stats::var_with_mean_bias(self, mean.clone(), dim);
        (var, mean)
    }

    /// Returns the median value along the specified dimension.
    ///
    /// The median is not unique for input tensors with an even number of elements
    /// in the reduced dimension. In this case, the lower of the two medians is returned,
    /// following PyTorch's behavior.
    ///
    /// # Note
    ///
    /// The current implementation performs a full sort along the specified dimension,
    /// which has O(nlog(n)) complexity. Additionally, most backends currently fall back
    /// to CPU for the sort operation, which may result in slower performance compared
    /// to native GPU operations.
    ///
    /// # Arguments
    ///
    /// - `dim` - The dimension along which to compute the median.
    ///   Negative dimensions are supported and count from the end.
    ///
    /// # Returns
    ///
    /// - A tensor containing the median values along the specified dimension.
    ///
    /// # Example 1
    ///
    /// ```ignore
    /// let device = Default::default();
    /// let tensor = Tensor::<2>::from_data(
    ///     [[1.0, 5.0, 3.0, 2.0], [8.0, 4.0, 6.0, 7.0]],
    ///     &device,
    /// );
    ///
    /// // Median along dimension 0:
    /// // sorted columns are [1.0, 8.0], [4.0, 5.0], [3.0, 6.0], [2.0, 7.0]
    /// let median = tensor.median(0);
    /// // Result: [[1.0, 4.0, 3.0, 2.0]]
    ///
    /// // Median along dimension 1:
    /// // sorted rows are [1.0, 2.0, 3.0, 5.0] and [4.0, 6.0, 7.0, 8.0]
    /// let median = tensor.median(1);
    /// // Result: [[2.0], [6.0]]
    /// ```
    ///
    /// # Example 2
    ///
    /// The median across all elements can be calculated as follows:
    ///
    /// ```ignore
    /// // D is the number of dimensions of the tensor
    /// let flattened_tensor: Tensor<1> = tensor.flatten(0, D - 1);
    ///
    /// // Calculate median for dim 0 since the tensor has become 1 dimensional
    /// let median = flattened_tensor.median(0);
    /// // Result: [4.0]
    /// ```
    pub fn median<I: AsIndex>(self, dim: I) -> Self {
        let dim = unwrap_dim_index(dim.try_dim_index(D), "Median");
        // TODO: Allow backend specialization. Optimally, implement a median kernel for cubecl
        // instead of leveraging a full sort to get the median.
        stats::median(self, dim)
    }

    /// Returns the median value along the specified dimension and its index.
    ///
    /// The median is not unique for input tensors with an even number of elements
    /// in the reduced dimension. In this case, the lower of the two medians is returned,
    /// following PyTorch's behavior.
    ///
    /// # Note
    ///
    /// The current implementation performs a full sort along the specified dimension,
    /// which has O(nlog(n)) complexity. Additionally, most backends currently fall back
    /// to CPU for the sort operation, which may result in slower performance compared
    /// to native GPU operations.
    ///
    /// # Arguments
    ///
    /// - `dim` - The dimension along which to compute the median.
    ///   Negative dimensions are supported and count from the end.
    ///
    /// # Returns
    ///
    /// A tuple containing:
    /// - A tensor with the median values.
    /// - A tensor with the indices of the median values in the original tensor.
    ///
    /// # Example
    ///
    /// ```ignore
    /// let device = Default::default();
    /// let tensor = Tensor::<2>::from_data(
    ///     [[1.0, 5.0, 3.0, 2.0], [8.0, 4.0, 6.0, 7.0]],
    ///     &device,
    /// );
    ///
    /// // Median along dimension 1:
    /// // sorted rows are [1.0, 2.0, 3.0, 5.0] and [4.0, 6.0, 7.0, 8.0]
    /// let (values, indices) = tensor.median_with_indices(1);
    /// // values: [[2.0], [6.0]], indices: [[3], [2]] (position in the original tensor)
    /// ```
    pub fn median_with_indices<I: AsIndex>(self, dim: I) -> (Self, Tensor<D, Int>) {
        let dim = unwrap_dim_index(dim.try_dim_index(D), "Median With Indices");
        // TODO: Allow backend specialization. Optimally, implement a median kernel for cubecl
        // instead of leveraging a full sort to get the median.
        stats::median_with_indices(self, dim)
    }

    /// Converts a tensor to the specified data type.
    ///
    /// Supports both within-kind casting (e.g., `FloatDType::F64`) and cross-kind casting
    /// (e.g., `IntDType::I64` to produce an int tensor).
    ///
    /// This is a no-op when casting to the current dtype within the same kind.
    ///
    /// # Example
    ///
    /// ```rust
    /// use burn_tensor::{Tensor, FloatDType, IntDType};
    ///
    /// let device = Default::default();
    /// let float_tensor = Tensor::<1>::from_floats([1.0, 2.5], &device);
    ///
    /// // Within-kind cast (float to float)
    /// let f64_tensor = float_tensor.clone().cast(FloatDType::F64);
    ///
    /// // Cross-kind cast (float to int)
    /// let int_tensor = float_tensor.cast(IntDType::I64);
    /// ```
    #[must_use]
    pub fn cast<T: Cast<D, Float>>(self, dtype: T) -> Tensor<D, T::OutputKind> {
        T::cast(self, dtype)
    }

    /// Detach the current tensor from the autodiff graph.
    ///
    /// This function does nothing when autodiff is not enabled.
    /// This can be used in batchers or elsewhere to ensure that previous operations are not
    /// considered in the autodiff graph.
    pub fn detach(self) -> Self {
        Self::new(detach_impl(self.primitive))
    }

    /// Mark the tensor to keep gradients during the backward pass.
    ///
    /// This function does nothing when autodiff is not enabled.
    pub fn require_grad(self) -> Self {
        self.set_require_grad(true)
    }

    /// Returns true if the tensor requires gradients during the backward pass.
    pub fn is_require_grad(&self) -> bool {
        is_require_grad_impl(&self.primitive)
    }

    /// Mark the tensor as tracked or untracked depending on the require_grad argument.
    /// When tracked, the gradients will be available after the backward pass.
    ///
    /// This function does nothing when autodiff is not enabled.
    pub fn set_require_grad(self, require_grad: bool) -> Self {
        Self::new(set_require_grad_impl(self.primitive, require_grad))
    }

    /// Applies the relu function to the tensor.
    pub(crate) fn relu(self) -> Self {
        Self::new(relu_impl(self.primitive))
    }

    /// Calculate covaraince matrix between different entries alongside a given dimension.
    ///
    /// # Arguments
    ///
    /// * `dim` - The dimension along which to calculate the covariance.
    ///   Negative dimensions are supported and count from the end.
    /// * `correction_factor` - Is usually 1 for samples and 0 for population.
    pub fn cov<I: AsIndex>(self, dim: I, correction_factor: usize) -> Tensor<D> {
        let dim = unwrap_dim_index(dim.try_dim_index(D), "Cov");
        let n = self.dims()[dim];
        let centered = (self.clone() - self.mean_dim(dim)).swap_dims(dim, 0);
        centered
            .clone()
            .transpose()
            .matmul(centered)
            .div_scalar(n as f32 - correction_factor as f32)
    }

    /// Convert the tensor to a lower precision data type based on the quantization scheme.
    ///
    /// # Arguments
    ///
    /// * `scheme` - The quantization scheme.
    /// * `qparams` - The pre-computed quantization parameters.
    ///
    /// # Returns
    ///
    /// The quantized tensor.
    pub fn quantize(self, scheme: &QuantScheme, qparams: QuantizationParameters) -> Tensor<D> {
        Tensor::new(quantize_impl(
            self.primitive,
            scheme,
            qparams.scales.primitive,
        ))
    }

    /// Dynamically convert the tensor to a lower precision data type based on the quantization scheme.
    ///
    /// # Arguments
    ///
    /// * `scheme` - The quantization scheme.
    ///
    /// # Returns
    ///
    /// The quantized tensor.
    ///
    /// # Notes
    /// This uses [min-max calibration](crate::quantization::Calibration::MinMax).
    pub fn quantize_dynamic(self, scheme: &QuantScheme) -> Tensor<D> {
        Tensor::new(quantize_dynamic_impl(self.primitive, scheme))
    }

    /// Convert the tensor back to a higher precision data type.
    ///
    /// If the tensor is not quantized, its value is simply returned.
    ///
    /// # Returns
    ///
    /// The dequantized tensor.
    pub fn dequantize(self) -> Tensor<D> {
        Tensor::new(dequantize_impl(self.primitive))
    }

    /// Checks element wise if the tensor is close to another tensor.
    ///
    /// The tolerance is defined by the following equation:
    ///
    /// ```text
    /// abs(a - b) <= (atol + rtol * abs(b))
    ///
    /// where `a` is the first tensor, `b` is the second tensor, `rtol` is the relative tolerance,
    /// and `atol` is the absolute tolerance.
    /// ```
    ///
    /// # Arguments
    ///
    /// * `other` - The tensor to compare with.
    /// * `rtol` - Optional relative tolerance. Default is 1e-5; see `DEFAULT_RTOL`.
    /// * `atol` - Optional absolute tolerance. Default is 1e-8; see `DEFAULT_ATOL`.
    ///
    /// # Returns
    ///
    /// A boolean tensor with the same shape as the input tensors.
    ///
    /// # Example
    ///
    /// ```rust
    /// use burn_tensor::{Tensor, Shape};
    ///
    /// let device = Default::default();
    /// let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
    /// let tensor2 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
    /// let tensor = tensor1.is_close(tensor2, None, None);
    /// println!("{tensor}");
    /// // [[true, true, true], [true, true, true]]
    /// ```
    pub fn is_close(self, other: Self, rtol: Option<f64>, atol: Option<f64>) -> Tensor<D, Bool> {
        let rtol = rtol.unwrap_or(DEFAULT_RTOL);
        let atol = atol.unwrap_or(DEFAULT_ATOL);

        // check finite difference is close
        let is_close_finite_val = self
            .clone()
            .sub(other.clone())
            .abs()
            .lower_equal(other.clone().abs().mul_scalar(rtol).add_scalar(atol))
            .bool_and(self.clone().is_finite())
            .bool_and(other.clone().is_finite());

        // check if both are infinite and have same sign
        let inf_same_sign = self
            .clone()
            .is_finite()
            .bool_not()
            .bool_and(other.clone().is_finite().bool_not())
            .bool_and(self.equal(other));

        is_close_finite_val.bool_or(inf_same_sign)
    }

    /// Checks if all elements are close to another tensor.
    ///
    /// The tolerance is defined by the following equation:
    ///
    /// ```text
    ///
    /// abs(a - b) <= (atol + rtol * abs(b))
    ///
    /// where `a` is the first tensor, `b` is the second tensor, `rtol` is the relative tolerance,
    /// and `atol` is the absolute tolerance.
    ///
    /// ```
    ///
    /// # Arguments
    ///
    /// * `other` - The tensor to compare with.
    /// * `rtol` - Optional relative tolerance. Default is 1e-5; see `DEFAULT_RTOL`.
    /// * `atol` - Optional absolute tolerance. Default is 1e-8; see `DEFAULT_ATOL`.
    ///
    /// # Returns
    ///
    /// A boolean scalar.
    ///
    /// # Remarks
    ///
    /// # Example
    ///
    /// ```rust
    /// use burn_tensor::{Tensor, Shape};
    ///
    /// let device = Default::default();
    /// let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
    /// let tensor2 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
    /// let result = tensor1.all_close(tensor2, None, None);
    /// println!("{}", result);
    /// // true
    /// ```
    pub fn all_close(self, other: Self, rtol: Option<f64>, atol: Option<f64>) -> bool {
        self.is_close(other, rtol, atol)
            .all()
            .into_scalar::<u8>()
            .to_bool()
    }

    /// Returns a new tensor with boolean elements indicating whether each element of the input is NaN.
    ///
    /// # Returns
    ///
    /// A boolean tensor where `true` indicates NaN and `false` indicates a non-NaN value.
    ///
    /// # Example
    ///
    /// ```rust
    /// use burn_tensor::{Tensor, Bool, Shape};
    ///
    /// let device = Default::default();
    /// let tensor = Tensor::<2>::from_data([[1.0, f64::NAN, 3.0], [5.0, 9.0, 6.0]], &device);
    /// let tensor = tensor.is_nan();
    /// println!("{tensor}");
    /// // [[false, true, false], [false, false, false]]
    /// ```
    pub fn is_nan(self) -> Tensor<D, Bool> {
        Tensor::new(is_nan_impl(self.primitive))
    }

    /// Checks if the tensor contains any NaN values.
    ///
    /// # Returns
    ///
    /// A boolean tensor with a single element indicating whether the tensor contains any NaN values.
    ///
    /// # Example
    ///
    /// ```rust
    /// use burn_tensor::{Tensor, Bool, Shape};
    ///
    /// let device = Default::default();
    /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [f64::NAN, 9.0, 6.0]], &device);
    /// let tensor = tensor.contains_nan();
    /// println!("{tensor}");
    /// // [true]
    /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
    /// let tensor = tensor.contains_nan();
    /// println!("{tensor}");
    /// // [false]
    /// ```
    pub fn contains_nan(self) -> Tensor<1, Bool> {
        self.is_nan().any()
    }

    /// Returns a new tensor with boolean elements indicating whether each element of the input is infinite (either +INF or -INF).
    ///
    /// # Returns
    ///
    /// A boolean tensor where `true` indicates that the value is infinite
    ///
    /// # Example
    ///
    /// ```rust
    /// use burn_tensor::{Tensor, Bool, Shape};
    ///
    /// let device = Default::default();
    /// let tensor = Tensor::<2>::from_data([[1.0, f64::INFINITY, 3.0], [f64::NAN, 9.0, 6.0]], &device);
    /// let tensor = tensor.is_finite();
    /// println!("{tensor}");
    /// // [[false, true, false], [false, false, false]]
    /// ```
    pub fn is_inf(self) -> Tensor<D, Bool> {
        Tensor::new(is_inf_impl(self.primitive))
    }

    /// Returns a new tensor with boolean elements indicating whether each element of the input is finite
    ///
    /// # Returns
    ///
    /// A boolean tensor where `true` indicates that the value is finite and `false` indicates
    /// either INF, -INF or NAN
    ///
    /// # Example
    ///
    /// ```rust
    /// use burn_tensor::{Tensor, Bool, Shape};
    ///
    /// let device = Default::default();
    /// let tensor = Tensor::<2>::from_data([[1.0, f64::INFINITY, 3.0], [f64::NAN, 9.0, 6.0]], &device);
    /// let tensor = tensor.is_finite();
    /// println!("{tensor}");
    /// // [[true, false, true], [false, true, true]]
    /// ```
    pub fn is_finite(self) -> Tensor<D, Bool> {
        self.clone()
            .is_nan()
            .bool_not()
            .bool_and(self.is_inf().bool_not())
    }

    /// Samples tensor as a two-dimensional spatial grid of (possibly multi-channel) values,
    /// using the given locations in [-1, 1].
    ///
    /// # Arguments
    ///
    /// * `grid` - A tensor of locations, with shape (N, H_out, W_out, 2). Values are [-1, 1].
    ///   A [x = -1, y = -1] means top-left, and [x = 1, y = 1] means bottom-right
    /// * `options` - Grid sampling options (mode, padding_mode, align_corners)
    ///
    /// # Returns
    ///
    /// A tensor with shape (N, C, H_out, W_out)
    ///
    /// # Example
    ///
    /// ```ignore
    /// use burn_tensor::ops::{GridSampleOptions, GridSamplePaddingMode, InterpolateMode};
    ///
    /// // Default options (bilinear, zeros padding, align_corners=false)
    /// let output = tensor.grid_sample_2d(grid, GridSampleOptions::default());
    ///
    /// // Custom options
    /// let options = GridSampleOptions::new(InterpolateMode::Bilinear)
    ///     .with_padding_mode(GridSamplePaddingMode::Border)
    ///     .with_align_corners(true);
    /// let output = tensor.grid_sample_2d(grid, options);
    /// ```
    pub fn grid_sample_2d(
        self,
        grid: Tensor<D>,
        options: impl Into<GridSampleOptions>,
    ) -> Tensor<D> {
        Tensor::new(grid_sample_2d_impl(
            self.primitive,
            grid.primitive,
            options.into(),
        ))
    }

    /// Computes the cross product of `self` and another tensor along a given dimension.
    ///
    /// Both `self` and `other` **must have size 3** along the specified `dim`,
    /// because the cross product is only defined in three-dimensional space.
    ///
    /// # Arguments
    ///
    /// * `other` - The other tensor to take the cross product with.
    /// * `dim`   - The dimension along which to compute the cross product.
    ///
    /// # Returns
    ///
    /// A tensor containing the cross product of `self` and `other` along `dim`.
    pub fn cross<Dim: AsIndex>(self, other: Tensor<D>, dim: Dim) -> Tensor<D> {
        let dim = unwrap_dim_index(dim.try_dim_index(D), "Cross");
        check!(TensorCheck::cross(&self, &other, dim));
        Tensor::new(cross_impl(self.primitive, other.primitive, dim))
    }

    /// Applies element wise power operation with a float Tensor
    ///
    /// # Arguments
    ///
    /// * `other` - The tensor to apply the power operation with.
    ///
    /// # Example
    ///
    /// ```rust
    /// use burn_tensor::{Tensor, Shape};
    ///
    /// let device = Default::default();
    /// let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
    /// let tensor2 = Tensor::<2>::from_data([[2.0, 3.0, 4.0], [1.0, 2.0, 3.0]], &device);
    /// let tensor = tensor1.powf(tensor2);
    /// println!("{tensor}");
    /// // [[1.0, 8.0, 81.0], [5.0, 81.0, 216.0]]
    /// ```
    pub fn powf(self, other: Self) -> Self {
        Tensor::new(powf_impl(self.primitive, other.primitive))
    }

    /// Applies element wise power operation with a float scalar
    ///
    /// # Arguments
    ///
    /// * `other` - The scalar to apply the power operation with.
    ///
    /// # Example
    ///
    /// ```rust
    /// use burn_tensor::{Tensor, Shape};
    ///
    /// let device = Default::default();
    /// let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
    /// let tensor = tensor.powf_scalar(2.0);
    /// println!("{tensor}");
    /// // [[1.0, 4.0, 9.0], [25.0, 81.0, 36.0]]
    /// ```
    pub fn powf_scalar<E: ElementConversion>(self, other: E) -> Self {
        let rhs = Scalar::new(other, &self.dtype());
        Tensor::new(powf_scalar_impl(self.primitive, rhs))
    }
}

impl<const D: usize> Tensor<D> {
    /// Draws samples from a categorical distribution defined by the last dimension
    /// of the input tensor.
    ///
    /// The last dimension is treated as a (possibly unnormalized) set of weights
    /// defining a categorical distribution over categories. All leading dimensions
    /// are treated as batch dimensions. The method returns integer indices of the
    /// sampled categories.
    ///
    /// # Arguments
    ///
    /// * `num_samples` - Number of samples to draw per distribution. Must be >= 1.
    ///
    /// # Panics
    ///
    /// Panics if `num_samples` is 0.
    ///
    /// # Note
    ///
    /// Distributions with all-zero weights produce undefined (NaN-based) sampling
    /// results. Callers should ensure each distribution has at least one positive
    /// weight.
    ///
    /// # Returns
    ///
    /// An integer tensor with the same shape as the input, except the last dimension
    /// is replaced by `num_samples`, containing sampled category indices in
    /// `[0, num_categories)`.
    ///
    /// # Example
    ///
    /// ```rust
    /// use burn_tensor::Tensor;
    ///
    /// let device = Default::default();
    /// let probs = Tensor::<2>::from_floats(
    ///     [[0.0, 1.0, 0.0], [0.0, 0.0, 1.0]],
    ///     &device,
    /// );
    /// let samples = probs.categorical(4);
    /// // First row always samples index 1, second row always samples index 2
    /// println!("{samples}");
    /// ```
    pub fn categorical(self, num_samples: usize) -> Tensor<D, Int> {
        assert!(num_samples > 0, "categorical: num_samples must be >= 1");

        let shape = self.shape();
        let num_categories = shape[D - 1];
        let batch_size = (shape.num_elements() / num_categories).max(1);
        let device = self.device();
        let dtype = self.dtype();

        // Flatten leading dimensions into a single batch dimension: [batch, categories]
        let flat: Tensor<2> = self.reshape([batch_size, num_categories]);

        // Normalize weights to probabilities
        let sum = flat.clone().sum_dim(1); // [batch, 1]
        let probs = flat / sum;

        // Cumulative sum along categories dimension
        let cumsum = probs.cumsum(1); // [batch, categories]

        // Uniform random values for each sample
        let uniform = Tensor::<2>::random(
            [batch_size, num_samples],
            Distribution::Uniform(0.0, 1.0),
            (&device, dtype),
        ); // [batch, num_samples]

        // Expand dimensions for broadcasting:
        //   cumsum: [batch, categories, 1]
        //   uniform: [batch, 1, num_samples]
        let cumsum_3d: Tensor<3> = cumsum.unsqueeze_dim(2);
        let uniform_3d: Tensor<3> = uniform.unsqueeze_dim(1);

        // Count categories where cumsum < uniform (inverse CDF)
        let mask: Tensor<3, Bool> = cumsum_3d.lower(uniform_3d);
        let indices: Tensor<2, Int> = mask.int().sum_dim(1).squeeze_dim::<2>(1);

        // Clamp to valid range to guard against floating-point imprecision in cumsum
        let indices = indices.clamp(0, num_categories as i64 - 1);

        // Reshape back to [...leading_dims, num_samples]
        let mut out_shape = shape;
        out_shape[D - 1] = num_samples;
        indices.reshape(out_shape)
    }
}

#[cfg(feature = "std")]
impl<const D: usize> Tensor<D> {
    /// Returns true if the tensor is marked as distributed.
    pub fn is_distributed(&self) -> bool {
        is_distributed_impl(&self.primitive)
    }

    /// Mark the tensor as distributed.
    ///
    /// This function does nothing when autodiff or distributed is not enabled.
    pub fn set_distributed(self, param_id: DistributedParamId) -> Self {
        Self::new(set_distributed_impl(self.primitive, param_id))
    }
}

impl<const D: usize, K> Tensor<D, K>
where
    K: FloatMath,
{
    /// Applies element wise square operation.
    ///
    #[cfg_attr(doc, doc = r#"$y_i = x_i * x_i$"#)]
    #[cfg_attr(not(doc), doc = "`y_i = x_i * x_i`")]
    pub fn square(self) -> Self {
        Self::new(K::square(self.primitive))
    }

    /// Applies element wise exponential operation.
    ///
    #[cfg_attr(doc, doc = r#"$y_i = e^{x_i}$"#)]
    #[cfg_attr(not(doc), doc = "`y = e^x`")]
    pub fn exp(self) -> Self {
        Self::new(K::exp(self.primitive))
    }

    /// Applies element wise natural logarithm of one plus the input tensor.
    ///
    #[cfg_attr(doc, doc = r#"$y_i = \log_e\(x_i + 1\)$"#)]
    #[cfg_attr(not(doc), doc = "`y_i = log1p(x_i)`")]
    pub fn log1p(self) -> Self {
        Self::new(K::log1p(self.primitive))
    }

    /// Applies element wise natural log operation *ln*.
    ///
    #[cfg_attr(doc, doc = r#"$y_i = \log_e\(x_i\)$"#)]
    #[cfg_attr(not(doc), doc = "`y_i = log(x_i)`")]
    pub fn log(self) -> Self {
        Self::new(K::log(self.primitive))
    }

    /// Applies element wise square root operation.
    ///
    pub fn sqrt(self) -> Self {
        Tensor::new(K::sqrt(self.primitive))
    }
    /// Applies element wise cosine operation.
    ///
    #[cfg_attr(doc, doc = r#"$y_i = \cos\(x_i\)$"#)]
    #[cfg_attr(not(doc), doc = "`y_i = cos(x_i)`")]
    pub fn cos(self) -> Self {
        Tensor::new(K::cos(self.primitive))
    }

    /// Applies element wise sine operation.
    ///
    #[cfg_attr(doc, doc = r#"$y_i = \sin\(x_i\)$"#)]
    #[cfg_attr(not(doc), doc = "`y_i = sin(x_i)`")]
    pub fn sin(self) -> Self {
        Tensor::new(K::sin(self.primitive))
    }

    /// Applies element wise tangent operation.
    ///
    #[cfg_attr(doc, doc = r#"$y_i = \tan\(x_i\)$"#)]
    #[cfg_attr(not(doc), doc = "`y_i = tan(x_i)`")]
    pub fn tan(self) -> Self {
        Tensor::new(K::tan(self.primitive))
    }

    /// Applies element wise hyperbolic cosine operation.
    ///
    #[cfg_attr(doc, doc = r#"$y_i = \cosh\(x_i\)$"#)]
    #[cfg_attr(not(doc), doc = "`y_i = cosh(x_i)`")]
    ///
    /// # Example
    ///
    /// ```rust
    /// use burn_tensor::Tensor;
    ///
    /// let device = Default::default();
    ///
    /// let tensor = Tensor::<1>::from_data([0.0, -1.0, 2.0], &device);
    /// println!("{}", tensor.cosh()); // [1.0, 1.5430, 3.7621]
    /// ```
    pub fn cosh(self) -> Self {
        Tensor::new(K::cosh(self.primitive))
    }

    /// Applies element wise hyperbolic sine operation.
    ///
    #[cfg_attr(doc, doc = r#"$y_i = \sinh\(x_i\)$"#)]
    #[cfg_attr(not(doc), doc = "`y_i = sinh(x_i)`")]
    ///
    /// # Example
    ///
    /// ```rust
    /// use burn_tensor::Tensor;
    ///
    /// let device = Default::default();
    ///
    /// let tensor = Tensor::<1>::from_data([0.0, -1.0, 2.0], &device);
    /// println!("{}", tensor.sinh()); // [0.0, -1.1752, 3.6269]
    /// ```
    pub fn sinh(self) -> Self {
        Tensor::new(K::sinh(self.primitive))
    }

    /// Applies element wise hyperbolic tangent operation.
    ///
    #[cfg_attr(doc, doc = r#"$y_i = \tanh\(x_i\)$"#)]
    #[cfg_attr(not(doc), doc = "`y_i = tanh(x_i)`")]
    ///
    /// # Example
    ///
    /// ```rust
    /// use burn_tensor::Tensor;
    ///
    /// let device = Default::default();
    ///
    /// let tensor = Tensor::<1>::from_data([0.0, -1.0, 2.0], &device);
    /// println!("{}", tensor.tanh()); // [0.0, -0.7616, 0.9640]
    /// ```
    pub fn tanh(self) -> Self {
        Tensor::new(K::tanh(self.primitive))
    }

    /// Applies element wise inverse cosine operation.
    ///
    #[cfg_attr(doc, doc = r#"$y_i = \acos\(x_i\)$"#)]
    #[cfg_attr(not(doc), doc = "`y_i = acos(x_i)`")]
    ///
    /// # Example
    ///
    /// ```rust
    /// use burn_tensor::Tensor;
    ///
    /// let device = Default::default();
    ///
    /// let tensor = Tensor::<1>::from_data([0.0, -1.0, 1.0], &device);
    /// println!("{}", tensor.acos()); // [1.5708, 3.1416, 0.0]
    /// ```
    pub fn acos(self) -> Self {
        Tensor::new(K::acos(self.primitive))
    }

    /// Applies element wise inverse hyperbolic cosine operation.
    ///
    #[cfg_attr(doc, doc = r#"$y_i = \acosh\(x_i\)$"#)]
    #[cfg_attr(not(doc), doc = "`y_i = acosh(x_i)`")]
    ///
    /// # Example
    ///
    /// ```rust
    /// use burn_tensor::Tensor;
    ///
    /// let device = Default::default();
    ///
    /// let tensor = Tensor::<1>::from_data([1.0, 2.0, 3.0], &device);
    /// println!("{}", tensor.acosh()); // [0.0000, 1.3170, 1.7627]
    /// ```
    pub fn acosh(self) -> Self {
        Tensor::new(K::acosh(self.primitive))
    }

    /// Applies element wise inverse sine operation.
    ///
    #[cfg_attr(doc, doc = r#"$y_i = \asin\(x_i\)$"#)]
    #[cfg_attr(not(doc), doc = "`y_i = asin(x_i)`")]
    ///
    /// # Example
    ///
    /// ```rust
    /// use burn_tensor::Tensor;
    ///
    /// let device = Default::default();
    ///
    /// let tensor = Tensor::<1>::from_data([0.0, -1.0, 1.0], &device);
    /// println!("{}", tensor.asin()); // [ 0.0000, -1.5708,  1.5708]
    /// ```
    pub fn asin(self) -> Self {
        Tensor::new(K::asin(self.primitive))
    }

    /// Applies element wise inverse hyperbolic sine operation.
    ///
    #[cfg_attr(doc, doc = r#"$y_i = \asinh\(x_i\)$"#)]
    #[cfg_attr(not(doc), doc = "`y_i = asinh(x_i)`")]
    ///
    /// # Example
    ///
    /// ```rust
    /// use burn_tensor::Tensor;
    ///
    /// let device = Default::default();
    ///
    /// let tensor = Tensor::<1>::from_data([0.0, -1.0, 1.0], &device);
    /// println!("{}", tensor.asinh()); // [ 0.0000, -0.8814,  0.8814]
    /// ```
    pub fn asinh(self) -> Self {
        Tensor::new(K::asinh(self.primitive))
    }

    /// Applies element wise inverse tangent operation.
    ///
    #[cfg_attr(doc, doc = r#"$y_i = \atan\(x_i\)$"#)]
    #[cfg_attr(not(doc), doc = "`y_i = atan(x_i)`")]
    ///
    /// # Example
    ///
    /// ```rust
    /// use burn_tensor::Tensor;
    ///
    /// let device = Default::default();
    ///
    /// let tensor = Tensor::<1>::from_data([0.0, -1.0, 2.0], &device);
    /// println!("{}", tensor.atan()); // [ 0.0, -0.7854,  1.1071]
    /// ```
    pub fn atan(self) -> Self {
        Tensor::new(K::atan(self.primitive))
    }

    /// Applies element wise inverse hyperbolic tangent operation.
    ///
    #[cfg_attr(doc, doc = r#"$y_i = \atanh\(x_i\)$"#)]
    #[cfg_attr(not(doc), doc = "`y_i = atanh(x_i)`")]
    ///
    /// # Example
    ///
    /// ```rust
    /// use burn_tensor::Tensor;
    ///
    /// let device = Default::default();
    ///
    /// let tensor = Tensor::<1>::from_data([0.0, -0.5, 0.5], &device);
    /// println!("{}", tensor.atanh()); // [ 0.0, -0.5493,  0.5493]
    /// ```
    pub fn atanh(self) -> Self {
        Tensor::new(K::atanh(self.primitive))
    }

    /// Applies element wise inverse tangent operation using the signs of arguments to determine the correct quadrant.
    ///
    #[cfg_attr(doc, doc = r#"$z_i = \atan2\(y_i, x_i\)$"#)]
    #[cfg_attr(not(doc), doc = "`z_i = atan2(y_i, x_i)`")]
    ///
    /// # Example
    ///
    /// ```rust
    /// use burn_tensor::Tensor;
    ///
    /// let device = Default::default();
    ///
    /// let lhs = Tensor::<1>::from_data([-2.0, 2.0, -2.0], &device);
    /// let rhs = Tensor::<1>::from_data([1.0, -1.0, -1.0], &device);
    /// println!("{}", lhs.atan2(rhs)); // [-1.1071,  2.0344, -2.0344]
    /// ```
    pub fn atan2(self, other: Self) -> Self {
        Tensor::new(K::atan2(self.primitive, other.primitive))
    }
}

// =========================================================================
// Non-generic implementation helpers (outlined from the public generic API).
// See the crate-level docs for the rationale behind this pattern.
// =========================================================================

fn erf_impl(p: BridgeTensor) -> BridgeTensor {
    BridgeTensor::float(Dispatch::float_erf(p.into_float()))
}

fn recip_impl(p: BridgeTensor) -> BridgeTensor {
    BridgeTensor::float(Dispatch::float_recip(p.into_float()))
}

fn hypot_impl(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor {
    BridgeTensor::float(Dispatch::float_hypot(lhs.into_float(), rhs.into_float()))
}

fn round_impl(p: BridgeTensor) -> BridgeTensor {
    BridgeTensor::float(Dispatch::float_round(p.into_float()))
}

fn floor_impl(p: BridgeTensor) -> BridgeTensor {
    BridgeTensor::float(Dispatch::float_floor(p.into_float()))
}

fn ceil_impl(p: BridgeTensor) -> BridgeTensor {
    BridgeTensor::float(Dispatch::float_ceil(p.into_float()))
}

fn int_impl(p: BridgeTensor, device: Device) -> BridgeTensor {
    let out_dtype = device.settings().int_dtype;
    BridgeTensor::int(Dispatch::float_into_int(p.into_float(), out_dtype))
}

fn random_like_impl(p: &BridgeTensor, distribution: Distribution) -> BridgeTensor {
    BridgeTensor::float(Dispatch::float_random(
        p.shape(),
        distribution,
        &p.as_dispatch().device(),
        p.dtype().into(),
    ))
}

fn detach_impl(p: BridgeTensor) -> BridgeTensor {
    BridgeTensor::float(Dispatch::float_detach(p.into_float()))
}

fn is_require_grad_impl(p: &BridgeTensor) -> bool {
    let (kind, tensor) = p.as_parts();
    match kind {
        BridgeKind::Float => Dispatch::float_is_require_grad(tensor),
        BridgeKind::QFloat => Dispatch::q_is_require_grad(tensor),
        _ => panic!("Should be Float primitive kind"),
    }
}

fn set_require_grad_impl(p: BridgeTensor, require_grad: bool) -> BridgeTensor {
    let (kind, tensor) = p.into_parts();
    match kind {
        BridgeKind::Float => {
            BridgeTensor::float(Dispatch::float_set_require_grad(tensor, require_grad))
        }
        BridgeKind::QFloat => {
            BridgeTensor::qfloat(Dispatch::q_set_require_grad(tensor, require_grad))
        }
        _ => panic!("Should be Float primitive kind"),
    }
}

fn relu_impl(p: BridgeTensor) -> BridgeTensor {
    BridgeTensor::float(Dispatch::relu(p.into_float()))
}

fn quantize_impl(p: BridgeTensor, scheme: &QuantScheme, scales: BridgeTensor) -> BridgeTensor {
    BridgeTensor::qfloat(Dispatch::quantize(
        p.into_float(),
        scheme,
        QuantizationParametersPrimitive {
            scales: scales.into_float(),
        },
    ))
}

fn quantize_dynamic_impl(p: BridgeTensor, scheme: &QuantScheme) -> BridgeTensor {
    BridgeTensor::qfloat(Dispatch::quantize_dynamic(p.into_float(), scheme))
}

fn dequantize_impl(p: BridgeTensor) -> BridgeTensor {
    BridgeTensor::float(p.into_float())
}

fn is_nan_impl(p: BridgeTensor) -> BridgeTensor {
    let bool_dtype = p.device_settings().bool_dtype;
    BridgeTensor::bool(Dispatch::float_is_nan(p.into_float(), bool_dtype))
}

fn is_inf_impl(p: BridgeTensor) -> BridgeTensor {
    let bool_dtype = p.device_settings().bool_dtype;
    BridgeTensor::bool(Dispatch::float_is_inf(p.into_float(), bool_dtype))
}

fn grid_sample_2d_impl(
    p: BridgeTensor,
    grid: BridgeTensor,
    options: GridSampleOptions,
) -> BridgeTensor {
    BridgeTensor::float(Dispatch::float_grid_sample_2d(
        p.into_float(),
        grid.into_float(),
        options,
    ))
}

fn cross_impl(p: BridgeTensor, other: BridgeTensor, dim: usize) -> BridgeTensor {
    BridgeTensor::float(Dispatch::float_cross(
        p.into_float(),
        other.into_float(),
        dim,
    ))
}

fn powf_impl(lhs: BridgeTensor, rhs: BridgeTensor) -> BridgeTensor {
    let (lkind, lhs) = lhs.into_parts();
    let (rkind, rhs) = rhs.into_parts();
    match (lkind, rkind) {
        (BridgeKind::Float, BridgeKind::Float) => {
            BridgeTensor::float(Dispatch::float_powf(lhs, rhs))
        }
        (BridgeKind::QFloat, BridgeKind::QFloat) => match Dispatch::q_powf(lhs, rhs) {
            TensorPrimitive::Float(out) => BridgeTensor::float(out),
            TensorPrimitive::QFloat(out) => BridgeTensor::qfloat(out),
        },
        (BridgeKind::QFloat, BridgeKind::Float) => {
            let dtype = rhs.dtype();
            BridgeTensor::float(Dispatch::float_powf(
                Dispatch::dequantize(lhs, dtype.into()),
                rhs,
            ))
        }
        (BridgeKind::Float, BridgeKind::QFloat) => {
            let dtype = lhs.dtype();
            BridgeTensor::float(Dispatch::float_powf(
                lhs,
                Dispatch::dequantize(rhs, dtype.into()),
            ))
        }
        _ => panic!("Should be Float primitive kind"),
    }
}

fn powf_scalar_impl(p: BridgeTensor, rhs: Scalar) -> BridgeTensor {
    let (kind, lhs) = p.into_parts();
    match kind {
        BridgeKind::Float => BridgeTensor::float(Dispatch::float_powf_scalar(lhs, rhs)),
        BridgeKind::QFloat => match Dispatch::q_powf_scalar(lhs, rhs) {
            TensorPrimitive::Float(out) => BridgeTensor::float(out),
            TensorPrimitive::QFloat(out) => BridgeTensor::qfloat(out),
        },
        _ => panic!("Should be Float primitive kind"),
    }
}

#[cfg(feature = "std")]
fn is_distributed_impl(p: &BridgeTensor) -> bool {
    let (kind, tensor) = p.as_parts();
    match kind {
        BridgeKind::Float => Dispatch::is_distributed(tensor),
        BridgeKind::QFloat => unimplemented!(),
        _ => panic!("Should be Float primitive kind"),
    }
}

#[cfg(feature = "std")]
fn set_distributed_impl(p: BridgeTensor, param_id: DistributedParamId) -> BridgeTensor {
    let (kind, tensor) = p.into_parts();
    match kind {
        BridgeKind::Float => {
            BridgeTensor::float(Dispatch::set_distributed_params(tensor, param_id))
        }
        BridgeKind::QFloat => unimplemented!(),
        _ => panic!("Should be Float primitive kind"),
    }
}