[Feat] Add new burn-vision crate with one initial op

Cargo.toml

@@ -153,8 +153,8 @@ ahash = { version = "0.8.11", default-features = false }
 portable-atomic-util = { version = "0.2.4", features = ["alloc"] }
 
 ### For the main burn branch. ###
-cubecl = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "a172f6760052bef392e6f0e44e912460960f2c1b" }
-cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "a172f6760052bef392e6f0e44e912460960f2c1b" }
+cubecl = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "ff94be8c9a79d1ac2e44829c2b4ec5a7e91b82e2" }
+cubecl-common = { git = "https://github.com/tracel-ai/cubecl", default-features = false, rev = "ff94be8c9a79d1ac2e44829c2b4ec5a7e91b82e2" }
 ### For local development. ###
 # cubecl = { path = "../cubecl/crates/cubecl", default-features = false }
 # cubecl-common = { path = "../cubecl/crates/cubecl-common", default-features = false }

crates/burn-candle/src/element.rs

@@ -4,8 +4,11 @@ use burn_tensor::Element;
 use candle_core::{FloatDType, Tensor, WithDType};
 use half::{bf16, f16};
 
+/// Candle element
 pub trait CandleElement: Element + WithDType {}
+/// Candle float element
 pub trait FloatCandleElement: CandleElement + FloatDType {}
+/// Candle int element
 pub trait IntCandleElement: CandleElement {}
 
 impl CandleElement for f64 {}

crates/burn-candle/src/lib.rs

@@ -13,6 +13,7 @@ mod ops;
 mod tensor;
 
 pub use backend::*;
+pub use element::*;
 pub use tensor::*;
 
 #[cfg(test)]

crates/burn-cuda/Cargo.toml

@@ -12,8 +12,8 @@ repository = "https://github.com/tracel-ai/burn/tree/main/crates/burn-cuda"
 version.workspace = true
 
 [features]
-default = ["fusion", "autotune", "burn-jit/default", "cubecl/default"]
 autotune = ["burn-jit/autotune"]
+default = ["fusion", "autotune", "burn-jit/default", "cubecl/default"]
 doc = ["burn-jit/doc"]
 fusion = ["burn-fusion", "burn-jit/fusion"]
 std = ["burn-jit/std", "cubecl/std"]

crates/burn-jit/src/kernel/index/mod.rs

@@ -11,7 +11,7 @@ pub(crate) use flip::*;
 pub(crate) use repeat_dim::*;
 pub(crate) use select::*;
 pub(crate) use select_assign::*;
-pub(crate) use slice::*;
+pub use slice::*;
 pub(crate) use slice_assign::*;
 
 pub(crate) use gather::*;

crates/burn-jit/src/kernel/index/slice.rs

@@ -3,7 +3,8 @@ use burn_tensor::Shape;
 use cubecl::{calculate_cube_count_elemwise, prelude::*};
 use std::ops::Range;
 
-pub(crate) fn slice<R: JitRuntime, E: JitElement>(
+/// Slice a jit tensor with a set of ranges
+pub fn slice<R: JitRuntime, E: JitElement>(
     tensor: JitTensor<R>,
     indices: &[Range<usize>],
 ) -> JitTensor<R> {

crates/burn-jit/src/kernel/mod.rs

@@ -39,4 +39,4 @@ pub mod reduce;
 
 pub(crate) use clamp::*;
 pub(crate) use comparison::*;
-pub(crate) use index::*;
+pub use index::*;

crates/burn-jit/src/lib.rs

@@ -7,7 +7,8 @@
 extern crate derive_new;
 extern crate alloc;
 
-mod ops;
+/// Utilities for implementing JIT kernels
+pub mod ops;
 
 /// Kernel module
 pub mod kernel;

crates/burn-jit/src/ops/base.rs

@@ -76,6 +76,7 @@ pub(crate) fn swap_dims<R: JitRuntime>(
     tensor
 }
 
+/// Permute a tensor's dimensions
 pub fn permute<R: JitRuntime>(mut tensor: JitTensor<R>, axes: &[usize]) -> JitTensor<R> {
     // remap strides
     tensor.strides = axes.iter().map(|i| tensor.strides[*i]).collect();
@@ -135,7 +136,8 @@ pub(crate) fn expand<R: JitRuntime>(tensor: JitTensor<R>, target_shape: Shape) -
     }
 }
 
-pub(crate) fn reshape<R: JitRuntime>(tensor: JitTensor<R>, shape: Shape) -> JitTensor<R> {
+/// Reshape a jit tensor to a new shape
+pub fn reshape<R: JitRuntime>(tensor: JitTensor<R>, shape: Shape) -> JitTensor<R> {
     // TODO: Not force standard layout all the time (improve performance).
     let tensor = kernel::into_contiguous(tensor);
 

crates/burn-jit/src/ops/mod.rs

@@ -7,6 +7,7 @@ mod qtensor;
 mod transaction;
 
 pub(crate) mod base;
-pub(crate) use base::*;
+pub use base::*;
 
-pub(crate) mod numeric;
+/// Numeric utility functions for jit backends
+pub mod numeric;