Replace some powf->powi (#3152)

ArthurBrussee · web-flow · commit 7e547c45fca8 · 2025-05-06T08:21:21.000-04:00
* Replace some powf-&gt;powi

* Fix
diff --git a/crates/burn-autodiff/src/ops/tensor.rs b/crates/burn-autodiff/src/ops/tensor.rs
@@ -375,15 +375,16 @@ impl<B: Backend, C: CheckpointStrategy> FloatTensorOps<Self> for Autodiff<B, C>
                     grads,
                     |grad| {
                         let rhs = rhs_4lhs.unwrap();
-                        let value = B::float_powf_scalar(rhs, -1.0);
+                        let value = B::float_recip(rhs);
                         let grad = B::float_mul(grad, value);
 
                         broadcast.backward_lhs::<B>(grad)
                     },
                     |grad| {
                         let rhs = rhs_4rhs.unwrap();
                         let lhs = lhs.unwrap();
-                        let value = B::float_div(B::float_neg(lhs), B::float_powf_scalar(rhs, 2.0));
+                        let value =
+                            B::float_div(B::float_neg(lhs), B::float_powi_scalar(rhs, 2.elem()));
                         let grad = B::float_mul(grad, value);
 
                         broadcast.backward_rhs::<B>(grad)
@@ -644,7 +645,7 @@ impl<B: Backend, C: CheckpointStrategy> FloatTensorOps<Self> for Autodiff<B, C>
             ) {
                 let tensor = checkpointer.retrieve_node_output(ops.state);
                 unary::<B, _>(ops.parents, ops.node, grads, |grad| {
-                    let tmp = B::float_powf_scalar(tensor, -2.0);
+                    let tmp = B::float_powi_scalar(tensor, (-2).elem());
                     let value = B::float_neg(tmp);
 
                     B::float_mul(grad, value)
@@ -1631,7 +1632,7 @@ impl<B: Backend, C: CheckpointStrategy> FloatTensorOps<Self> for Autodiff<B, C>
             ) {
                 let input = checkpointer.retrieve_node_output(ops.state);
                 unary::<B, _>(ops.parents, ops.node, grads, |grad| {
-                    let value = B::float_powf_scalar(input, -1.0);
+                    let value = B::float_recip(input);
                     B::float_mul(grad, value)
                 });
             }
@@ -1670,7 +1671,7 @@ impl<B: Backend, C: CheckpointStrategy> FloatTensorOps<Self> for Autodiff<B, C>
                 let input = checkpointer.retrieve_node_output(ops.state);
                 unary::<B, _>(ops.parents, ops.node, grads, |grad| {
                     let value = B::float_add_scalar(input, 1.elem());
-                    let value = B::float_powf_scalar(value, -1.0);
+                    let value = B::float_recip(value);
 
                     B::float_mul(grad, value)
                 });
@@ -1920,7 +1921,7 @@ impl<B: Backend, C: CheckpointStrategy> FloatTensorOps<Self> for Autodiff<B, C>
                 let state = B::float_tanh(input);
                 unary::<B, _>(ops.parents, ops.node, grads, |grad| {
                     let value = B::float_add_scalar(
-                        B::float_neg(B::float_powf_scalar(state, 2.0)),
+                        B::float_neg(B::float_powi_scalar(state, 2.elem())),
                         1.elem(),
                     );
                     B::float_mul(grad, value)
@@ -2068,7 +2069,7 @@ impl<B: Backend, C: CheckpointStrategy> FloatTensorOps<Self> for Autodiff<B, C>
             ) {
                 unary::<B, _>(ops.parents, ops.node, grads, |grad| {
                     let ops = checkpointer.retrieve_node_output(ops.state);
-                    let exponent = B::float_neg(B::float_powf_scalar(ops, 2.0));
+                    let exponent = B::float_neg(B::float_powi_scalar(ops, 2.elem()));
                     let numerator = B::float_mul_scalar(B::float_exp(exponent), 2.0.elem());
                     let denominator = core::f64::consts::PI.sqrt().elem();
                     let value = B::float_div_scalar(numerator, denominator);
diff --git a/crates/burn-core/src/grad_clipping/base.rs b/crates/burn-core/src/grad_clipping/base.rs
@@ -88,7 +88,7 @@ impl GradientClipping {
     }
 
     fn l2_norm<B: Backend, const D: usize>(tensor: Tensor<B, D>) -> Tensor<B, 1> {
-        let squared = tensor.powf_scalar(2.0);
+        let squared = tensor.powi_scalar(2);
         let sum = squared.sum();
         sum.sqrt()
     }
diff --git a/crates/burn-core/src/nn/loss/huber.rs b/crates/burn-core/src/nn/loss/huber.rs
@@ -132,7 +132,7 @@ impl HuberLoss {
         // Moreover |r| = sign(r) * r
         let outside = softsign.mul(residuals.clone()).sub_scalar(self.lin_bias);
 
-        let inside = residuals.powf_scalar(2.).mul_scalar(0.5);
+        let inside = residuals.powi_scalar(2).mul_scalar(0.5);
         inside.mask_where(is_large, outside)
     }
 }
diff --git a/crates/burn-core/src/nn/loss/mse.rs b/crates/burn-core/src/nn/loss/mse.rs
@@ -46,7 +46,7 @@ impl MseLoss {
         logits: Tensor<B, D>,
         targets: Tensor<B, D>,
     ) -> Tensor<B, D> {
-        logits.sub(targets).powf_scalar(2.0)
+        logits.sub(targets).powi_scalar(2)
     }
 }
 
diff --git a/crates/burn-core/src/nn/norm/batch.rs b/crates/burn-core/src/nn/norm/batch.rs
@@ -138,7 +138,7 @@ impl<const D: usize, B: Backend> BatchNorm<B, D> {
         let var = input
             .clone()
             .sub(mean.clone())
-            .powf_scalar(2.0)
+            .powi_scalar(2)
             .swap_dims(0, 1)
             .reshape([channels, flatten_size])
             .mean_dim(1)
diff --git a/crates/burn-core/src/nn/norm/group.rs b/crates/burn-core/src/nn/norm/group.rs
@@ -170,7 +170,7 @@ pub(crate) fn group_norm<B: Backend, const D: usize>(
     let mean = input.clone().sum_dim(2) / hidden_size as f64;
     let input = input.sub(mean);
 
-    let var = input.clone().powf_scalar(2.).sum_dim(2) / hidden_size as f64;
+    let var = input.clone().powi_scalar(2).sum_dim(2) / hidden_size as f64;
     let input_normalized = input.div(var.add_scalar(epsilon).sqrt());
 
     if affine {
diff --git a/crates/burn-core/src/nn/norm/rms.rs b/crates/burn-core/src/nn/norm/rms.rs
@@ -71,8 +71,7 @@ impl<B: Backend> RmsNorm<B> {
     pub fn forward<const D: usize>(&self, x: Tensor<B, D>) -> Tensor<B, D> {
         // Calculate the root-mean-square norm of the input tensor along the last dimension
         let dtype = x.dtype();
-        let rms =
-            (x.clone().cast(DType::F32).powf_scalar(2.0).mean_dim(D - 1) + self.epsilon).sqrt();
+        let rms = (x.clone().cast(DType::F32).powi_scalar(2).mean_dim(D - 1) + self.epsilon).sqrt();
         (x / rms.cast(dtype)) * self.gamma.val().unsqueeze()
     }
 }
diff --git a/crates/burn-core/src/nn/rope_encoding.rs b/crates/burn-core/src/nn/rope_encoding.rs
@@ -79,7 +79,7 @@ impl RotaryEncodingConfig {
         // Calculate (10000 ^ (2i / d_model)) by using the log base property `exp(log(10000) * (2i / d_model))`
         // This is done since burn doesn't support exponentiation of scalar to tensor
         let theta_i = exponent.mul_scalar(self.theta.ln()).exp();
-        let theta_i = theta_i.powf_scalar(-1.0);
+        let theta_i = theta_i.recip();
 
         let theta_i = scaling(theta_i);
 
diff --git a/crates/burn-core/src/optim/adagrad.rs b/crates/burn-core/src/optim/adagrad.rs
@@ -117,11 +117,11 @@ impl LrDecay {
         lr_decay_state: Option<LrDecayState<B, D>>,
     ) -> (Tensor<B, D>, LrDecayState<B, D>) {
         let state = if let Some(mut state) = lr_decay_state {
-            state.sum = state.sum.add(grad.clone().powf_scalar(2.));
+            state.sum = state.sum.add(grad.clone().powi_scalar(2));
             state.time += 1;
             state
         } else {
-            LrDecayState::new(1, grad.clone().powf_scalar(2.))
+            LrDecayState::new(1, grad.clone().powi_scalar(2))
         };
 
         let new_lr = lr / (1. + (state.time as f64 - 1.) * self.lr_decay);
diff --git a/crates/burn-core/src/optim/adam.rs b/crates/burn-core/src/optim/adam.rs
@@ -10,7 +10,7 @@ use super::{
 use crate::config::Config;
 use crate::optim::adaptor::OptimizerAdaptor;
 use crate::tensor::{Tensor, backend::AutodiffBackend};
-use burn_tensor::{ElementConversion, backend::Backend, ops::Device};
+use burn_tensor::{backend::Backend, ops::Device};
 
 #[cfg(not(feature = "std"))]
 use num_traits::Float;
@@ -140,7 +140,7 @@ impl AdaptiveMomentum {
             state.moment_2 = state
                 .moment_2
                 .mul_scalar(self.beta_2)
-                .add(grad.powf_scalar(2.0).mul_scalar(factor));
+                .add(grad.powi_scalar(2).mul_scalar(factor));
 
             state.time += 1;
 
@@ -150,12 +150,12 @@ impl AdaptiveMomentum {
             let moment_1 = grad.clone().mul_scalar(factor);
 
             let factor = 1.0 - self.beta_2;
-            let moment_2 = grad.powf_scalar(2.0).mul_scalar(factor);
+            let moment_2 = grad.powi_scalar(2).mul_scalar(factor);
 
             AdaptiveMomentumState::new(1, moment_1, moment_2)
         };
 
-        let time = (state.time as i32).elem();
+        let time = state.time as i32;
         let moment_1_corrected = state
             .moment_1
             .clone()
diff --git a/crates/burn-core/src/optim/adamw.rs b/crates/burn-core/src/optim/adamw.rs
@@ -126,7 +126,7 @@ impl AdaptiveMomentumW {
             state.moment_2 = state
                 .moment_2
                 .mul_scalar(self.beta_2)
-                .add(grad.powf_scalar(2.0).mul_scalar(factor));
+                .add(grad.powi_scalar(2).mul_scalar(factor));
 
             // Update time.
             state.time += 1;
@@ -139,7 +139,7 @@ impl AdaptiveMomentumW {
 
             // Initialize second moment estimate.
             let factor = 1.0 - self.beta_2;
-            let moment_2 = grad.powf_scalar(2.0).mul_scalar(factor);
+            let moment_2 = grad.powi_scalar(2).mul_scalar(factor);
 
             AdaptiveMomentumState::new(1, moment_1, moment_2)
         };
diff --git a/crates/burn-core/src/optim/rmsprop.rs b/crates/burn-core/src/optim/rmsprop.rs
@@ -159,11 +159,11 @@ impl<B: Backend, const D: usize> SquareAvgState<B, D> {
                 let square_avg = state
                     .square_avg
                     .mul_scalar(alpha)
-                    .add(grad.clone().powf_scalar(2.).mul_scalar(1. - alpha));
+                    .add(grad.clone().powi_scalar(2).mul_scalar(1. - alpha));
                 (grad, Self { square_avg })
             }
             _ => {
-                let square_avg = grad.clone().powf_scalar(2.).mul_scalar(1. - alpha);
+                let square_avg = grad.clone().powi_scalar(2).mul_scalar(1. - alpha);
                 (grad, Self { square_avg })
             }
         }
@@ -215,7 +215,7 @@ impl<B: Backend, const D: usize> CenteredState<B, D> {
             let avg = square_avg_state
                 .square_avg
                 .clone()
-                .sub(grad_avg.clone().powf_scalar(2.));
+                .sub(grad_avg.clone().powi_scalar(2));
 
             (
                 grad,
diff --git a/crates/burn-tensor/src/tensor/api/numeric.rs b/crates/burn-tensor/src/tensor/api/numeric.rs
@@ -3801,7 +3801,7 @@ impl<B: Backend> Numeric<B> for Int {
     }
 
     fn powi_scalar<E: ElementConversion>(lhs: Self::Primitive, rhs: E) -> Self::Primitive {
-        B::int_powf_scalar(lhs, rhs.elem())
+        B::int_powi_scalar(lhs, rhs.elem())
     }
 
     fn random(shape: Shape, distribution: Distribution, device: &Device<B>) -> Self::Primitive {
@@ -4291,10 +4291,10 @@ impl<B: Backend> Numeric<B> for Float {
     fn powi_scalar<E: ElementConversion>(lhs: Self::Primitive, rhs: E) -> Self::Primitive {
         match lhs {
             TensorPrimitive::Float(lhs) => {
-                TensorPrimitive::Float(B::float_powf_scalar(lhs, rhs.elem()))
+                TensorPrimitive::Float(B::float_powi_scalar(lhs, rhs.elem()))
             }
             TensorPrimitive::QFloat(lhs) => {
-                TensorPrimitive::QFloat(B::q_powf_scalar(lhs, rhs.elem()))
+                TensorPrimitive::QFloat(B::q_powi_scalar(lhs, rhs.elem()))
             }
         }
     }
diff --git a/crates/burn-tensor/src/tensor/ops/activation.rs b/crates/burn-tensor/src/tensor/ops/activation.rs
@@ -101,7 +101,7 @@ pub trait ActivationOps<B: Backend> {
         let constant_3 = 0.0535161;
         let constant_4 = 0.398942;
 
-        let x3 = B::float_powf_scalar(x.clone(), 3.0);
+        let x3 = B::float_powi_scalar(x.clone(), 3.elem());
 
         let c1 = B::float_mul_scalar(x3.clone(), constant_1.elem());
         let c2 = B::float_mul_scalar(x.clone(), constant_2.elem());
@@ -113,7 +113,7 @@ pub trait ActivationOps<B: Backend> {
 
         let tanh = B::float_tanh(inner1);
 
-        let sech = B::float_powf_scalar(tanh.clone(), 2.0);
+        let sech = B::float_powi_scalar(tanh.clone(), 2.elem());
         let sech = B::float_neg(sech);
         let sech = B::float_add_scalar(sech, 1.elem());
 
diff --git a/crates/burn-tensor/src/tensor/ops/int_tensor.rs b/crates/burn-tensor/src/tensor/ops/int_tensor.rs
@@ -1,7 +1,6 @@
 use super::cat::cat_with_slice_assign;
 use super::repeat_dim::repeat_with_slice_assign;
 use super::{BoolTensor, Device, FloatTensor, IntElem, IntTensor};
-use crate::cast::ToElement;
 use crate::{Distribution, ElementConversion, Int, TensorData, backend::Backend, tensor::Shape};
 use alloc::vec::Vec;
 use core::ops::Range;
@@ -419,10 +418,7 @@ pub trait IntTensorOps<B: Backend> {
     ///
     /// The elements of `lhs` raised to the power of the elements of `rhs`.
     fn int_powi(lhs: IntTensor<B>, rhs: IntTensor<B>) -> IntTensor<B> {
-        B::float_into_int(B::float_powf(
-            B::int_into_float(lhs),
-            B::int_into_float(rhs),
-        ))
+        B::float_into_int(B::float_powi(B::int_into_float(lhs), rhs))
     }
 
     /// Element-wise power with a floatTensor.
@@ -450,7 +446,7 @@ pub trait IntTensorOps<B: Backend> {
     ///
     /// The elements of `lhs` raised to the value of `rhs`.
     fn int_powi_scalar(lhs: IntTensor<B>, rhs: IntElem<B>) -> IntTensor<B> {
-        B::float_into_int(B::float_powf_scalar(B::int_into_float(lhs), rhs.to_f32()))
+        B::float_into_int(B::float_powi_scalar(B::int_into_float(lhs), rhs))
     }
 
     /// Element-wise power with a floatTensor.
diff --git a/crates/burn-tensor/src/tensor/ops/tensor.rs b/crates/burn-tensor/src/tensor/ops/tensor.rs
@@ -1,7 +1,6 @@
 use super::cat::cat_with_slice_assign;
 use super::repeat_dim::repeat_with_slice_assign;
 use super::{BoolTensor, Device, FloatElem, FloatTensor, IntElem, IntTensor};
-use crate::tensor::cast::ToElement;
 use crate::{Distribution, ElementConversion, Float, TensorData, backend::Backend, tensor::Shape};
 use crate::{FloatDType, TensorMetadata, TensorPrimitive};
 use alloc::vec::Vec;
@@ -845,7 +844,7 @@ pub trait FloatTensorOps<B: Backend> {
     ///
     /// The elements of `lhs` raised to the value of `rhs`.
     fn float_powi_scalar(lhs: FloatTensor<B>, rhs: IntElem<B>) -> FloatTensor<B> {
-        Self::float_powf_scalar(lhs, rhs.to_f32())
+        Self::float_powf_scalar(lhs, rhs.elem::<f32>())
     }
 
     /// Returns a new tensor with values raised to the power of float `value`.
diff --git a/crates/burn-tensor/src/tensor/stats/mod.rs b/crates/burn-tensor/src/tensor/stats/mod.rs
@@ -36,7 +36,7 @@ pub fn var_with_mean_n<B: Backend, const D: usize>(
 ) -> Tensor<B, D> {
     tensor
         .sub(mean)
-        .powf_scalar(2.0)
+        .powi_scalar(2)
         .sum_dim(dim)
         .div_scalar(n as f32)
 }

Original file line number	Diff line number	Diff line change
`@@ -88,7 +88,7 @@ impl GradientClipping {`
`88`	`88`	`}`
`89`	`89`
`90`	`90`	`fn l2_norm<B: Backend, const D: usize>(tensor: Tensor<B, D>) -> Tensor<B, 1> {`
`91`		`- let squared = tensor.powf_scalar(2.0);`
	`91`	`+ let squared = tensor.powi_scalar(2);`
`92`	`92`	`let sum = squared.sum();`
`93`	`93`	`sum.sqrt()`
`94`	`94`	`}`
Original file line number	Diff line number	Diff line change
`@@ -132,7 +132,7 @@ impl HuberLoss {`
`132`	`132`	`// Moreover \|r\| = sign(r) * r`
`133`	`133`	`let outside = softsign.mul(residuals.clone()).sub_scalar(self.lin_bias);`
`134`	`134`
`135`		`- let inside = residuals.powf_scalar(2.).mul_scalar(0.5);`
	`135`	`+ let inside = residuals.powi_scalar(2).mul_scalar(0.5);`
`136`	`136`	`inside.mask_where(is_large, outside)`
`137`	`137`	`}`
`138`	`138`	`}`
Original file line number	Diff line number	Diff line change
`@@ -46,7 +46,7 @@ impl MseLoss {`
`46`	`46`	`logits: Tensor<B, D>,`
`47`	`47`	`targets: Tensor<B, D>,`
`48`	`48`	`) -> Tensor<B, D> {`
`49`		`- logits.sub(targets).powf_scalar(2.0)`
	`49`	`+ logits.sub(targets).powi_scalar(2)`
`50`	`50`	`}`
`51`	`51`	`}`
`52`	`52`
Original file line number	Diff line number	Diff line change
`@@ -71,8 +71,7 @@ impl<B: Backend> RmsNorm<B> {`
`71`	`71`	`pub fn forward<const D: usize>(&self, x: Tensor<B, D>) -> Tensor<B, D> {`
`72`	`72`	`// Calculate the root-mean-square norm of the input tensor along the last dimension`
`73`	`73`	`let dtype = x.dtype();`
`74`		`- let rms =`
`75`		`- (x.clone().cast(DType::F32).powf_scalar(2.0).mean_dim(D - 1) + self.epsilon).sqrt();`
	`74`	`+ let rms = (x.clone().cast(DType::F32).powi_scalar(2).mean_dim(D - 1) + self.epsilon).sqrt();`
`76`	`75`	`(x / rms.cast(dtype)) * self.gamma.val().unsqueeze()`
`77`	`76`	`}`
`78`	`77`	`}`