Sparse tensors

crates/burn-core/Cargo.toml

@@ -71,6 +71,7 @@ vision = ["burn-dataset?/vision", "burn-common/network"]
 # Backend
 autodiff = ["burn-autodiff"]
 fusion = ["burn-wgpu?/fusion"]
+sparse = ["burn-sparse"]
 
 ## Backend features
 metal = ["burn-candle?/metal"]
@@ -116,6 +117,7 @@ burn-cuda = { path = "../burn-cuda", version = "0.14.0", optional = true, defaul
 burn-autodiff = { path = "../burn-autodiff", version = "0.14.0", optional = true }
 burn-tch = { path = "../burn-tch", version = "0.14.0", optional = true }
 burn-candle = { path = "../burn-candle", version = "0.14.0", optional = true }
+burn-sparse = { path = "../burn-sparse", version = "0.14.0", optional = true }
 
 derive-new = { workspace = true }
 log = { workspace = true, optional = true }

crates/burn-core/src/backend.rs

@@ -33,3 +33,6 @@ pub use burn_tch as libtorch;
 
 #[cfg(feature = "tch")]
 pub use burn_tch::LibTorch;
+
+#[cfg(feature = "sparse")]
+pub use burn_sparse as sparse;

crates/burn-sparse/Cargo.toml

@@ -0,0 +1,43 @@
+[package]
+authors = []
+categories = ["science", "no-std", "embedded", "wasm"]
+description = "Sparse tensor crate that offers a default sparse backend wrapper around burn backends."
+edition.workspace = true
+keywords = ["deep-learning", "machine-learning", "tensor", "sparse"]
+license.workspace = true
+name = "burn-sparse"
+readme.workspace = true
+repository = "https://github.com/tracel-ai/burn/tree/main/burn-sparse"
+version.workspace = true
+
+[features]
+default = ["std"]
+doc = ["default"]
+experimental-named-tensor = []
+std = ["rand/std", "half/std", "num-traits/std"]
+wasm-sync = []
+
+[dependencies]
+burn-common = { path = "../burn-common", version = "0.14.0", default-features = false }
+burn-tensor = { path = "../burn-tensor", version = "0.14.0" }
+
+proc-macro2 = { workspace = true }
+quote = { workspace = true }
+syn = { workspace = true }
+derive-new = { workspace = true }
+half = { workspace = true }
+num-traits = { workspace = true }
+rand = { workspace = true }
+rand_distr = { workspace = true } # use instead of statrs because it supports no_std
+
+# The same implementation of HashMap in std but with no_std support (only needs alloc crate)
+hashbrown = { workspace = true } # no_std compatible
+
+# Serialization
+serde = { workspace = true }
+
+[dev-dependencies]
+rand = { workspace = true, features = ["std", "std_rng"] } # Default enables std
+
+[package.metadata.docs.rs]
+features = ["doc"]

crates/burn-sparse/src/coo.rs

@@ -0,0 +1,76 @@
+use burn_tensor::backend::Backend;
+use burn_tensor::ops::SparseBoolOps;
+use burn_tensor::ops::SparseTensorOps;
+use burn_tensor::Dense;
+use burn_tensor::Device;
+use burn_tensor::Float;
+use burn_tensor::Int;
+use burn_tensor::Shape;
+use burn_tensor::Sparse;
+use burn_tensor::SparseStorage;
+use burn_tensor::Tensor;
+use burn_tensor::TensorData;
+use burn_tensor::TensorKind;
+
+#[derive(Clone, Debug)]
+pub struct COO;
+
+#[derive(Clone, Debug)]
+pub struct SparseCOOTensor<B: Backend, K: TensorKind<B>, const D: usize> {
+    pub coordinates: Option<Tensor<B, 2, Int>>,
+    pub values: Option<Tensor<B, 1, K>>,
+    pub shape: Shape<D>,
+    pub device: Device<B>,
+}
+
+impl<B: Backend> SparseStorage<B> for COO {
+    type SparsePrimitive<K: burn_tensor::TensorKind<B>, const D: usize> = SparseCOOTensor<B, K, D>;
+
+    fn name() -> &'static str {
+        "SparseCOO"
+    }
+}
+
+impl<B: Backend> SparseTensorOps<COO, B> for COO {}
+
+pub(crate) fn flatten_coordinates<B: Backend, const D: usize, const S: usize>(
+    coordinates: Tensor<B, 2, Int>,
+    shape: Shape<D>,
+    device: &Device<B>,
+) -> Tensor<B, 2, Int> {
+    let mut strides_data = [[1]; D];
+    for i in (0..D).rev() {
+        if D - 1 - i == S {
+            strides_data[i] = [1];
+        } else if D - 1 - i < S {
+            strides_data[i] = [0];
+        } else {
+            strides_data[i] = [strides_data[i + 1][0] * shape.dims[i + 1] as i64];
+        }
+    }
+    let strides_data: TensorData = TensorData::from(strides_data);
+    let strides: Tensor<B, 2, Int> = Tensor::from_data(strides_data, device);
+    let flat_coordinates: Tensor<B, 1, Int> = strides.mul(coordinates).sum_dim(0).flatten(0, 1);
+
+    flat_coordinates.unsqueeze_dim(0)
+}
+
+pub(crate) fn unflatten_coordinates<B: Backend, const D: usize>(
+    flat_coordinates: Tensor<B, 2, Int>,
+    new_shape: Shape<D>,
+) -> Tensor<B, 2, Int> {
+    let flat_coordinates = flat_coordinates.squeeze::<1>(0);
+    let mut remaining_flat_coordinates = flat_coordinates.clone();
+    let mut new_coordinates = Vec::with_capacity(D);
+
+    for &dim_size in new_shape.dims.iter().rev() {
+        let size = dim_size as i64;
+        let new_coord = remaining_flat_coordinates.clone().remainder_scalar(size);
+        new_coordinates.push(new_coord.clone());
+        remaining_flat_coordinates = remaining_flat_coordinates.div_scalar(size);
+    }
+
+    new_coordinates.reverse();
+
+    Tensor::stack(new_coordinates, 0)
+}

crates/burn-sparse/src/coo_bool.rs

@@ -0,0 +1,163 @@
+use super::coo::COO;
+use crate::SparseCOOTensor;
+use crate::{flatten_coordinates, unflatten_coordinates};
+use burn_tensor::Int;
+use burn_tensor::ReprPrimitive;
+use burn_tensor::Shape;
+use burn_tensor::Tensor;
+use burn_tensor::{
+    backend::Backend,
+    ops::{SparseBoolOps, SparseTensorOps},
+    SparseStorage,
+};
+use burn_tensor::{Bool, Dense};
+
+impl<B: Backend> SparseBoolOps<COO, B> for COO {
+    fn bool_to_sparse<const D: usize>(
+        dense: <B as Backend>::BoolTensorPrimitive<D>,
+    ) -> <COO as SparseStorage<B>>::SparsePrimitive<burn_tensor::Bool, D> {
+        todo!()
+    }
+
+    fn bool_empty<const D: usize>(
+        shape: burn_tensor::Shape<D>,
+        device: &burn_tensor::Device<B>,
+    ) -> <COO as SparseStorage<B>>::SparsePrimitive<burn_tensor::Bool, D> {
+        todo!()
+    }
+
+    fn bool_shape<const D: usize>(
+        tensor: &<COO as SparseStorage<B>>::SparsePrimitive<burn_tensor::Bool, D>,
+    ) -> burn_tensor::Shape<D> {
+        todo!()
+    }
+
+    fn bool_reshape<const D1: usize, const D2: usize>(
+        tensor: <COO as SparseStorage<B>>::SparsePrimitive<burn_tensor::Bool, D1>,
+        shape: burn_tensor::Shape<D2>,
+    ) -> <COO as SparseStorage<B>>::SparsePrimitive<burn_tensor::Bool, D2> {
+        todo!()
+    }
+
+    fn bool_transpose<const D: usize>(
+        tensor: <COO as SparseStorage<B>>::SparsePrimitive<burn_tensor::Bool, D>,
+    ) -> <COO as SparseStorage<B>>::SparsePrimitive<burn_tensor::Bool, D> {
+        todo!()
+    }
+
+    fn bool_swap_dims<const D: usize>(
+        tensor: <COO as SparseStorage<B>>::SparsePrimitive<burn_tensor::Bool, D>,
+        dim1: usize,
+        dim2: usize,
+    ) -> <COO as SparseStorage<B>>::SparsePrimitive<burn_tensor::Bool, D> {
+        todo!()
+    }
+
+    fn bool_permute<const D: usize>(
+        tensor: <COO as SparseStorage<B>>::SparsePrimitive<burn_tensor::Bool, D>,
+        axes: &[usize],
+    ) -> <COO as SparseStorage<B>>::SparsePrimitive<burn_tensor::Bool, D> {
+        todo!()
+    }
+
+    fn bool_flip<const D: usize>(
+        tensor: <COO as SparseStorage<B>>::SparsePrimitive<burn_tensor::Bool, D>,
+        axes: &[usize],
+    ) -> <COO as SparseStorage<B>>::SparsePrimitive<burn_tensor::Bool, D> {
+        todo!()
+    }
+
+    fn bool_slice<const D1: usize, const D2: usize>(
+        tensor: <COO as SparseStorage<B>>::SparsePrimitive<burn_tensor::Bool, D1>,
+        indices: [std::ops::Range<usize>; D2],
+    ) -> <COO as SparseStorage<B>>::SparsePrimitive<burn_tensor::Bool, D1> {
+        todo!()
+    }
+
+    fn bool_slice_assign<const D1: usize, const D2: usize>(
+        tensor: <COO as SparseStorage<B>>::SparsePrimitive<burn_tensor::Bool, D1>,
+        ranges: [std::ops::Range<usize>; D2],
+        value: <COO as SparseStorage<B>>::SparsePrimitive<burn_tensor::Bool, D1>,
+    ) -> <COO as SparseStorage<B>>::SparsePrimitive<burn_tensor::Bool, D1> {
+        todo!()
+    }
+
+    fn bool_device<const D: usize>(
+        tensor: &<COO as SparseStorage<B>>::SparsePrimitive<burn_tensor::Bool, D>,
+    ) -> burn_tensor::Device<B> {
+        todo!()
+    }
+
+    fn bool_to_device<const D: usize>(
+        tensor: <COO as SparseStorage<B>>::SparsePrimitive<burn_tensor::Bool, D>,
+        device: &burn_tensor::Device<B>,
+    ) -> <COO as SparseStorage<B>>::SparsePrimitive<burn_tensor::Bool, D> {
+        todo!()
+    }
+
+    fn bool_repeat_dim<const D: usize>(
+        tensor: <COO as SparseStorage<B>>::SparsePrimitive<burn_tensor::Bool, D>,
+        dim: usize,
+        times: usize,
+    ) -> <COO as SparseStorage<B>>::SparsePrimitive<burn_tensor::Bool, D> {
+        todo!()
+    }
+
+    fn bool_cat<const D: usize>(
+        tensors: Vec<<COO as SparseStorage<B>>::SparsePrimitive<burn_tensor::Bool, D>>,
+        dim: usize,
+    ) -> <COO as SparseStorage<B>>::SparsePrimitive<burn_tensor::Bool, D> {
+        todo!()
+    }
+
+    fn bool_any<const D: usize>(
+        tensor: <COO as SparseStorage<B>>::SparsePrimitive<burn_tensor::Bool, D>,
+    ) -> <COO as SparseStorage<B>>::SparsePrimitive<burn_tensor::Bool, 1> {
+        todo!()
+    }
+
+    fn bool_any_dim<const D: usize>(
+        tensor: <COO as SparseStorage<B>>::SparsePrimitive<burn_tensor::Bool, D>,
+        dim: usize,
+    ) -> <COO as SparseStorage<B>>::SparsePrimitive<burn_tensor::Bool, D> {
+        todo!()
+    }
+
+    fn bool_all<const D: usize>(
+        tensor: <COO as SparseStorage<B>>::SparsePrimitive<burn_tensor::Bool, D>,
+    ) -> <COO as SparseStorage<B>>::SparsePrimitive<burn_tensor::Bool, 1> {
+        todo!()
+    }
+
+    fn bool_all_dim<const D: usize>(
+        tensor: <COO as SparseStorage<B>>::SparsePrimitive<burn_tensor::Bool, D>,
+        dim: usize,
+    ) -> <COO as SparseStorage<B>>::SparsePrimitive<burn_tensor::Bool, D> {
+        todo!()
+    }
+
+    fn bool_expand<const D1: usize, const D2: usize>(
+        tensor: <COO as SparseStorage<B>>::SparsePrimitive<burn_tensor::Bool, D1>,
+        shape: burn_tensor::Shape<D2>,
+    ) -> <COO as SparseStorage<B>>::SparsePrimitive<burn_tensor::Bool, D2> {
+        todo!()
+    }
+
+    fn bool_coordinates<const D: usize>(
+        mut tensor: <COO as SparseStorage<B>>::SparsePrimitive<burn_tensor::Bool, D>,
+    ) -> Option<ReprPrimitive<B, Int, Dense, 2>> {
+        tensor.coordinates.map(|c| c.into_primitive())
+    }
+
+    fn bool_to_dense<const D: usize>(
+        sparse: <COO as SparseStorage<B>>::SparsePrimitive<burn_tensor::Bool, D>,
+    ) -> B::BoolTensorPrimitive<D> {
+        todo!()
+    }
+
+    fn bool_values<const D: usize>(
+        tensor: ReprPrimitive<B, Bool, burn_tensor::Sparse<B, COO>, D>,
+    ) -> Option<ReprPrimitive<B, Bool, Dense, 1>> {
+        tensor.values.map(|v| v.into_primitive())
+    }
+}