refactor/mul-ops (tracel-ai#67)

djdisodo · Nov 5, 2022 · ee61e84 · ee61e84
1 parent 2bdad6f
commit ee61e84
Show file tree

Hide file tree

Showing 25 changed files with 330 additions and 352 deletions.
diff --git a/.github/workflows/test-burn-dataset.yml b/.github/workflows/test-burn-dataset.yml
@@ -14,7 +14,7 @@ jobs:
         uses: actions-rs/toolchain@v1
         with:
           profile: minimal
-          toolchain: beta
+          toolchain: stable
           components: rustfmt, clippy
           override: true
 

diff --git a/.github/workflows/test-burn-tensor.yml b/.github/workflows/test-burn-tensor.yml
@@ -14,7 +14,7 @@ jobs:
         uses: actions-rs/toolchain@v1
         with:
           profile: minimal
-          toolchain: beta
+          toolchain: stable
           components: rustfmt, clippy
           override: true
 

diff --git a/.github/workflows/test-burn.yml b/.github/workflows/test-burn.yml
@@ -14,7 +14,7 @@ jobs:
         uses: actions-rs/toolchain@v1
         with:
           profile: minimal
-          toolchain: beta
+          toolchain: stable
           components: rustfmt, clippy
           override: true
 

diff --git a/burn-tensor/src/tensor/backend/autodiff/ops/aggregation.rs b/burn-tensor/src/tensor/backend/autodiff/ops/aggregation.rs
@@ -76,9 +76,9 @@ impl<B: Backend, const D: usize> UnaryOps<B::TensorPrimitive<D>, B::TensorPrimit
         let ones = B::ones(shape, B::device(&grad));
 
         let val = 1_f64 / shape.dims[dim] as f64;
-        let ones = ones.mul_scalar(&B::Elem::from_elem(val));
+        let ones = B::mul_scalar(&ones, &B::Elem::from_elem(val));
 
-        ones.mul(&grad)
+        B::mul(&ones, &grad)
     }
 }
 
@@ -94,7 +94,7 @@ impl<B: Backend, const D: usize> UnaryOps<B::TensorPrimitive<D>, B::TensorPrimit
         let grad = state.output.grad().sum_dim(dim);
         let ones = B::ones(shape, B::device(&grad));
 
-        ones.mul(&grad)
+        B::mul(&ones, &grad)
     }
 }
 

diff --git a/burn-tensor/src/tensor/backend/autodiff/ops/div.rs b/burn-tensor/src/tensor/backend/autodiff/ops/div.rs
@@ -13,14 +13,14 @@ register_ops!(
         let value = state.right.value();
         let value = value.ones().div(&value);
 
-        state.output.grad().mul(&value)
+        B::mul(&state.output.grad(), &value)
     },
     partial_right |state: &BinaryOpsNodeState<B::TensorPrimitive<D>, B::TensorPrimitive<D>, B::TensorPrimitive<D>>| {
         let value_left = state.left.value();
         let value_right = state.right.value();
-        let value = value_left.neg().div(&value_right.mul(&value_right));
+        let value = value_left.neg().div(&B::mul(&value_right, &value_right));
 
-        state.output.grad().mul(&value)
+        B::mul(&state.output.grad(), &value)
     },
 );
 
@@ -31,7 +31,7 @@ register_ops!(
         let value = state_recorded.input.value();
         let tmp = value.ones().div_scalar(state);
 
-        state_recorded.output.grad().mul(&tmp)
+        B::mul(&state_recorded.output.grad(), &tmp)
     },
 );
 

diff --git a/burn-tensor/src/tensor/backend/autodiff/ops/erf.rs b/burn-tensor/src/tensor/backend/autodiff/ops/erf.rs
@@ -13,10 +13,11 @@ register_ops!(
     partial |state: &UnaryOpsNodeState<B::TensorPrimitive<D>, B::TensorPrimitive<D>>|{
         let value = state.input.value();
         let exponent = value.powf(2.0.to_elem()).neg();
-        let numerator = exponent.exp().mul_scalar(&2.0.to_elem());
+        let numerator = B::mul_scalar(&exponent.exp(), &2.0.to_elem());
         let denominator = std::f64::consts::PI.sqrt().to_elem();
         let value = numerator.div_scalar(&denominator);
-        state.output.grad().mul(&value)
+
+        B::mul(&state.output.grad(), &value)
     },
 );
 

diff --git a/burn-tensor/src/tensor/backend/autodiff/ops/exp.rs b/burn-tensor/src/tensor/backend/autodiff/ops/exp.rs
@@ -10,7 +10,7 @@ register_ops!(
     ops UnaryOps,
     name ADTensorExpOps,
     partial |state: &UnaryOpsNodeState<B::TensorPrimitive<D>, B::TensorPrimitive<D>>|{
-        state.output.grad().mul(&state.output.value())
+        B::mul(&state.output.grad(), &state.output.value())
     },
 );
 

diff --git a/burn-tensor/src/tensor/backend/autodiff/ops/log.rs b/burn-tensor/src/tensor/backend/autodiff/ops/log.rs
@@ -12,7 +12,7 @@ register_ops!(
     partial |state: &UnaryOpsNodeState<B::TensorPrimitive<D>, B::TensorPrimitive<D>>|{
         let value = state.input.value();
         let value = value.ones().div(&value);
-        state.output.grad().mul(&value)
+        B::mul(&state.output.grad(), &value)
     },
 );
 

diff --git a/burn-tensor/src/tensor/backend/autodiff/ops/mod.rs b/burn-tensor/src/tensor/backend/autodiff/ops/mod.rs
@@ -13,7 +13,6 @@ mod map_comparison;
 mod mask;
 mod matmul;
 mod module;
-mod mul;
 mod neg;
 mod pow;
 mod precision;

diff --git a/burn-tensor/src/tensor/backend/autodiff/ops/mul.rs b/burn-tensor/src/tensor/backend/autodiff/ops/mul.rs
diff --git a/burn-tensor/src/tensor/backend/autodiff/ops/pow.rs b/burn-tensor/src/tensor/backend/autodiff/ops/pow.rs
@@ -14,11 +14,11 @@ register_ops!(
         value: &f32,
         state: &UnaryOpsNodeState<B::TensorPrimitive<D>, B::TensorPrimitive<D>>
     | {
-        let value = state.input
+        let value = B::mul_scalar(&state.input
             .value()
             .powf(value - 1.0)
-            .mul_scalar(&value.clone().to_elem());
-        state.output.grad().mul(&value)
+            , &value.clone().to_elem());
+        B::mul(&state.output.grad(), &value)
     },
 );