Move testgen to burn-jit

tracel-ai · Jan 28, 2025 · 9e65150 · 9e65150
1 parent d292f85
commit 9e65150
Show file tree

Hide file tree

Showing 8 changed files with 8 additions and 19 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/burn-cuda/Cargo.toml b/crates/burn-cuda/Cargo.toml
@@ -37,9 +37,6 @@ log = { workspace = true }
 burn-jit = { path = "../burn-jit", version = "0.17.0", default-features = false, features = [
   "export_tests",
 ] }
-burn-vision = { path = "../burn-vision", version = "0.17.0", default-features = false, features = [
-  "export_tests",
-] }
 paste = { workspace = true }
 
 

diff --git a/crates/burn-cuda/src/lib.rs b/crates/burn-cuda/src/lib.rs
@@ -21,5 +21,4 @@ mod tests {
 
     // TODO: Add tests for bf16
     burn_jit::testgen_all!([f16, f32], [i8, i16, i32, i64], [u8, u32]);
-    burn_vision::testgen_all!();
 }
diff --git a/crates/burn-jit/Cargo.toml b/crates/burn-jit/Cargo.toml
@@ -20,9 +20,10 @@ export_tests = [
     "serial_test",
     "burn-autodiff/export_tests",
     "burn-tensor/export_tests",
-    "burn-vision?/export_tests",
+    "burn-vision/export_tests",
     "burn-ndarray",
     "fusion",
+    "vision",
     "paste",
 ]
 fusion = ["burn-fusion"]

diff --git a/crates/burn-jit/src/tests/mod.rs b/crates/burn-jit/src/tests/mod.rs
@@ -33,6 +33,7 @@ pub use burn_autodiff;
 pub use burn_fusion;
 pub use burn_ndarray;
 pub use burn_tensor;
+pub use burn_vision;
 pub use serial_test;
 
 #[macro_export]
@@ -43,7 +44,10 @@ macro_rules! testgen_all {
     };
     ([$($float:ident),*], [$($int:ident),*], [$($bool:ident),*]) => {
         mod jit {
+            pub use $crate::tests::burn_vision;
+
             burn_jit::testgen_jit!([$($float),*], [$($int),*], [$($bool),*]);
+            burn_vision::testgen_all!();
 
             mod kernel {
                 use super::*;

diff --git a/crates/burn-vision/src/tests/mod.rs b/crates/burn-vision/src/tests/mod.rs
@@ -5,12 +5,6 @@ macro_rules! testgen_all {
     () => {
         use burn_tensor::{Bool, Float, Int};
 
-        pub type TestBackend = JitBackend<TestRuntime, f32, i32, u32>;
-
-        type TestTensor<const D: usize> = burn_tensor::Tensor<TestBackend, D>;
-        type TestTensorInt<const D: usize> = burn_tensor::Tensor<TestBackend, D, Int>;
-        type TestTensorBool<const D: usize> = burn_tensor::Tensor<TestBackend, D, Bool>;
-
         pub mod vision {
             pub use super::*;
 

diff --git a/crates/burn-wgpu/Cargo.toml b/crates/burn-wgpu/Cargo.toml
@@ -21,12 +21,12 @@ std = ["burn-jit/std", "cubecl/std"]
 template = ["burn-jit/template", "cubecl/template"]
 
 # Backends
-webgpu = ["cubecl-wgsl"]
 vulkan = ["cubecl-spirv"]
+webgpu = ["cubecl-wgsl"]
 
 # Compilers
-cubecl-wgsl = []
 cubecl-spirv = ["cubecl/wgpu-spirv"]
+cubecl-wgsl = []
 
 [dependencies]
 cubecl = { workspace = true, features = ["wgpu"] }
@@ -42,9 +42,6 @@ burn-tensor = { path = "../burn-tensor", version = "0.17.0", default-features =
 burn-jit = { path = "../burn-jit", version = "0.17.0", default-features = false, features = [
     "export_tests",
 ] }
-burn-vision = { path = "../burn-vision", version = "0.17.0", features = [
-    "export_tests",
-] }
 half = { workspace = true }
 paste = { workspace = true }
 

diff --git a/crates/burn-wgpu/src/lib.rs b/crates/burn-wgpu/src/lib.rs
@@ -124,5 +124,4 @@ mod tests {
     burn_jit::testgen_all!([f16, f32], [i8, i16, i32, i64], [u8, u32]);
     #[cfg(not(feature = "vulkan"))]
     burn_jit::testgen_all!([f32], [i32], [u32]);
-    burn_vision::testgen_all!();
 }