Make autodiff compile on wasm (#1889)

tracel-ai · Jun 14, 2024 · 47a8127 · 47a8127
1 parent 5e58ae1
commit 47a8127
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 3 deletions.
diff --git a/crates/burn-autodiff/src/ops/bool_tensor.rs b/crates/burn-autodiff/src/ops/bool_tensor.rs
@@ -121,10 +121,12 @@ impl<B: Backend, C: CheckpointStrategy> BoolTensorOps<Self> for Autodiff<B, C> {
         B::bool_flip(tensor, axes)
     }
 
+    #[cfg(any(feature = "wasm-sync", not(target_family = "wasm")))]
     fn bool_argwhere<const D: usize>(tensor: BoolTensor<B, D>) -> IntTensor<B, 2> {
         B::bool_argwhere(tensor)
     }
 
+    #[cfg(any(feature = "wasm-sync", not(target_family = "wasm")))]
     fn bool_nonzero<const D: usize>(tensor: BoolTensor<B, D>) -> Vec<IntTensor<B, 1>> {
         B::bool_nonzero(tensor)
     }

diff --git a/crates/burn-autodiff/src/ops/int_tensor.rs b/crates/burn-autodiff/src/ops/int_tensor.rs
@@ -383,6 +383,7 @@ impl<B: Backend, C: CheckpointStrategy> IntTensorOps<Self> for Autodiff<B, C> {
         B::int_expand(tensor, shape)
     }
 
+    #[cfg(any(feature = "wasm-sync", not(target_family = "wasm")))]
     fn int_sort<const D: usize>(
         tensor: IntTensor<Self, D>,
         dim: usize,
@@ -391,6 +392,7 @@ impl<B: Backend, C: CheckpointStrategy> IntTensorOps<Self> for Autodiff<B, C> {
         B::int_sort(tensor, dim, descending)
     }
 
+    #[cfg(any(feature = "wasm-sync", not(target_family = "wasm")))]
     fn int_sort_with_indices<const D: usize>(
         tensor: IntTensor<Self, D>,
         dim: usize,
@@ -399,6 +401,7 @@ impl<B: Backend, C: CheckpointStrategy> IntTensorOps<Self> for Autodiff<B, C> {
         B::int_sort_with_indices(tensor, dim, descending)
     }
 
+    #[cfg(any(feature = "wasm-sync", not(target_family = "wasm")))]
     fn int_argsort<const D: usize>(
         tensor: IntTensor<Self, D>,
         dim: usize,

diff --git a/crates/burn-autodiff/src/ops/tensor.rs b/crates/burn-autodiff/src/ops/tensor.rs
@@ -21,7 +21,6 @@ use burn_tensor::{
 };
 
 use super::maxmin::MaxMinDim;
-use super::sort::SortDim;
 
 impl<B: Backend, C: CheckpointStrategy> FloatTensorOps<Self> for Autodiff<B, C> {
     fn float_from_data<const D: usize>(
@@ -2369,12 +2368,13 @@ impl<B: Backend, C: CheckpointStrategy> FloatTensorOps<Self> for Autodiff<B, C>
         }
     }
 
+    #[cfg(any(feature = "wasm-sync", not(target_family = "wasm")))]
     fn float_sort<const D: usize>(
         tensor: FloatTensor<Self, D>,
         dim: usize,
         descending: bool,
     ) -> FloatTensor<Self, D> {
-        match SortDim
+        match super::sort::SortDim
             .prepare::<C>([tensor.node])
             .compute_bound()
             .stateful()
@@ -2391,12 +2391,13 @@ impl<B: Backend, C: CheckpointStrategy> FloatTensorOps<Self> for Autodiff<B, C>
         }
     }
 
+    #[cfg(any(feature = "wasm-sync", not(target_family = "wasm")))]
     fn float_sort_with_indices<const D: usize>(
         tensor: FloatTensor<Self, D>,
         dim: usize,
         descending: bool,
     ) -> (FloatTensor<Self, D>, IntTensor<B, D>) {
-        match SortDim
+        match super::sort::SortDim
             .prepare::<C>([tensor.node])
             .compute_bound()
             .stateful()
@@ -2419,6 +2420,7 @@ impl<B: Backend, C: CheckpointStrategy> FloatTensorOps<Self> for Autodiff<B, C>
         }
     }
 
+    #[cfg(any(feature = "wasm-sync", not(target_family = "wasm")))]
     fn float_argsort<const D: usize>(
         tensor: FloatTensor<Self, D>,
         dim: usize,