Make ndarray tensor public (tracel-ai#920)

hojjatabdollahi · Nov 1, 2023 · 64e58b4 · 64e58b4
1 parent 8742d31
commit 64e58b4
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/burn-ndarray/src/lib.rs b/burn-ndarray/src/lib.rs
@@ -23,7 +23,7 @@ mod tensor;
 pub use backend::*;
 pub use element::FloatNdArrayElement;
 pub(crate) use sharing::*;
-pub(crate) use tensor::*;
+pub use tensor::*;
 
 extern crate alloc;
 

diff --git a/burn-ndarray/src/tensor.rs b/burn-ndarray/src/tensor.rs
@@ -2,8 +2,10 @@ use burn_tensor::{Data, Shape};
 
 use ndarray::{ArcArray, Array, Dim, IxDyn};
 
+/// Tensor primitive used by the [ndarray backend](crate::NdArray).
 #[derive(new, Debug, Clone)]
 pub struct NdArrayTensor<E, const D: usize> {
+    /// Dynamic array that contains the data of type E.
     pub array: ArcArray<E, IxDyn>,
 }
 
@@ -101,6 +103,7 @@ impl<E, const D: usize> NdArrayTensor<E, D>
 where
     E: Default + Clone,
 {
+    /// Create a new [ndarray tensor](NdArrayTensor) from [data](Data).
     pub fn from_data(data: Data<E, D>) -> NdArrayTensor<E, D> {
         let shape = data.shape.clone();
         let to_array = |data: Data<E, D>| Array::from_iter(data.value).into_shared();