tenferro_ext_ndarray/
lib.rs

1//! Bridge between [ndarray](https://docs.rs/ndarray) arrays and tenferro tensors.
2//!
3//! This crate provides typed canonical conversion helpers between generic
4//! `ndarray::ArrayBase<S, D>` inputs and [`tenferro_tensor::Tensor<T>`], plus
5//! export back to owned `ndarray::ArrayD<T>`. An optional `frontend` feature
6//! adds a convenience conversion into `tenferro::Tensor`.
7//!
8//! The canonical bridge normalizes at the boundary:
9//!
10//! - `ndarray -> tenferro_tensor::Tensor<T>` first materializes a standard
11//!   row-major ndarray layout, then converts into tenferro's internal
12//!   column-major canonical tensor layout
13//! - `tenferro_tensor::Tensor<T> -> ndarray` always materializes a standard
14//!   row-major owned ndarray result
15//!
16//! This keeps the interop contract simple and avoids exposing a separate
17//! zero-copy expert mode with layout-dependent reshape semantics.
18//!
19//! # Examples
20//!
21//! ```ignore
22//! use ndarray::Array2;
23//! use tenferro_ext_ndarray::{ndarray_to_tensor, tensor_to_ndarray};
24//!
25//! let array = Array2::from_shape_vec((2, 2), vec![1.0_f64, 2.0, 3.0, 4.0]).unwrap();
26//! let tensor = ndarray_to_tensor(array);
27//! let roundtrip = tensor_to_ndarray(tensor);
28//! assert_eq!(roundtrip.shape(), &[2, 2]);
29//! ```
30
31use ndarray::{ArrayBase, ArrayD, Data, Dimension, IxDyn};
32use tenferro_algebra::Scalar;
33use tenferro_device::{Error, LogicalMemorySpace, Result};
34use tenferro_tensor::{MemoryOrder, Tensor};
35
36#[cfg(test)]
37mod tests;
38
39#[cfg(feature = "frontend")]
40mod frontend {
41    use num_complex::{Complex32, Complex64};
42
43    pub trait Sealed {}
44
45    impl Sealed for f32 {}
46    impl Sealed for f64 {}
47    impl Sealed for Complex32 {}
48    impl Sealed for Complex64 {}
49}
50
51fn shape_error(err: ndarray::ShapeError) -> Error {
52    Error::InvalidArgument(format!("ndarray layout conversion failed: {err}"))
53}
54
55fn ensure_main_memory(space: LogicalMemorySpace) -> Result<()> {
56    if space != LogicalMemorySpace::MainMemory {
57        return Err(Error::InvalidArgument(
58            "tenferro-ext-ndarray currently supports CPU/main-memory tensors only".into(),
59        ));
60    }
61    Ok(())
62}
63
64fn row_major_strides(dims: &[usize]) -> Vec<isize> {
65    let ndim = dims.len();
66    if ndim == 0 {
67        return vec![];
68    }
69
70    let mut strides = vec![0isize; ndim];
71    strides[ndim - 1] = 1;
72    for i in (0..ndim - 1).rev() {
73        strides[i] = strides[i + 1] * dims[i + 1] as isize;
74    }
75    strides
76}
77
78fn standard_row_major_owned<T, S, D>(array: ArrayBase<S, D>) -> ArrayD<T>
79where
80    T: Clone,
81    S: Data<Elem = T>,
82    D: Dimension,
83{
84    let array = array.into_dyn();
85    if array.is_standard_layout() {
86        array.into_owned()
87    } else {
88        array.as_standard_layout().to_owned()
89    }
90}
91
92fn into_owned_data<T: Scalar>(tensor: Tensor<T>, context: &str) -> Result<Vec<T>> {
93    tensor
94        .try_into_data_vec()
95        .ok_or_else(|| Error::InvalidArgument(context.into()))
96}
97
98/// Fallibly converts an ndarray array into a typed tenferro tensor.
99///
100/// This is the canonical interop entry point. The bridge first materializes a
101/// standard row-major ndarray layout, then normalizes it into tenferro's
102/// internal column-major tensor layout.
103///
104/// # Examples
105///
106/// ```ignore
107/// use ndarray::Array2;
108/// use tenferro_ext_ndarray::try_ndarray_to_tensor;
109///
110/// let array = Array2::from_shape_vec((1, 2), vec![1.0_f64, 2.0]).unwrap();
111/// let tensor = try_ndarray_to_tensor(array).unwrap();
112/// assert_eq!(tensor.dims(), &[1, 2]);
113/// ```
114pub fn try_ndarray_to_tensor<T, S, D>(array: ArrayBase<S, D>) -> Result<Tensor<T>>
115where
116    T: Scalar + Clone,
117    S: Data<Elem = T>,
118    D: Dimension,
119{
120    let array = standard_row_major_owned(array);
121    let dims = array.shape().to_vec();
122    let (data, offset) = array.into_raw_vec_and_offset();
123    let tensor = Tensor::from_vec(
124        data,
125        &dims,
126        &row_major_strides(&dims),
127        offset.unwrap_or(0) as isize,
128    )?;
129    Ok(tensor.into_contiguous(MemoryOrder::ColumnMajor))
130}
131
132/// Converts an owned ndarray array into a typed tenferro tensor, panicking on
133/// conversion failure.
134///
135/// # Examples
136///
137/// ```ignore
138/// use ndarray::Array2;
139/// use tenferro_ext_ndarray::ndarray_to_tensor;
140///
141/// let array = Array2::from_shape_vec((1, 2), vec![1.0_f64, 2.0]).unwrap();
142/// let tensor = ndarray_to_tensor(array);
143/// assert_eq!(tensor.dims(), &[1, 2]);
144/// ```
145pub fn ndarray_to_tensor<T, S, D>(array: ArrayBase<S, D>) -> Tensor<T>
146where
147    T: Scalar + Clone,
148    S: Data<Elem = T>,
149    D: Dimension,
150{
151    try_ndarray_to_tensor(array).unwrap_or_else(|err| panic!("{err}"))
152}
153
154/// Fallibly converts a typed tenferro tensor into an owned ndarray array.
155///
156/// The canonical bridge always materializes a standard row-major owned ndarray
157/// result. This keeps row-major downstream integrations simple while tenferro
158/// continues to use column-major canonical tensors internally.
159///
160/// # Examples
161///
162/// ```ignore
163/// use tenferro_device::LogicalMemorySpace;
164/// use tenferro_ext_ndarray::try_tensor_to_ndarray;
165/// use tenferro_tensor::{MemoryOrder, Tensor};
166///
167/// let tensor = Tensor::<f64>::zeros(&[2, 2], LogicalMemorySpace::MainMemory, MemoryOrder::RowMajor).unwrap();
168/// let array = try_tensor_to_ndarray(tensor).unwrap();
169/// assert_eq!(array.shape(), &[2, 2]);
170/// ```
171pub fn try_tensor_to_ndarray<T: Scalar>(tensor: Tensor<T>) -> Result<ArrayD<T>> {
172    ensure_main_memory(tensor.logical_memory_space())?;
173
174    let row_major = tensor.into_contiguous(MemoryOrder::RowMajor);
175    let dims = row_major.dims().to_vec();
176    let data = into_owned_data(
177        row_major,
178        "into_contiguous(RowMajor) must yield an owned CPU buffer for ndarray export",
179    )?;
180    ArrayD::from_shape_vec(IxDyn(&dims), data).map_err(shape_error)
181}
182
183/// Converts a typed tenferro tensor into an owned ndarray array, panicking on
184/// conversion failure.
185///
186/// # Examples
187///
188/// ```ignore
189/// use tenferro_device::LogicalMemorySpace;
190/// use tenferro_ext_ndarray::tensor_to_ndarray;
191/// use tenferro_tensor::{MemoryOrder, Tensor};
192///
193/// let tensor = Tensor::<f64>::zeros(&[2], LogicalMemorySpace::MainMemory, MemoryOrder::ColumnMajor).unwrap();
194/// let array = tensor_to_ndarray(tensor);
195/// assert_eq!(array.shape(), &[2]);
196/// ```
197pub fn tensor_to_ndarray<T: Scalar>(tensor: Tensor<T>) -> ArrayD<T> {
198    try_tensor_to_ndarray(tensor).unwrap_or_else(|err| panic!("{err}"))
199}
200
201/// Marker trait for scalar dtypes supported by the optional frontend helper.
202///
203/// This trait is sealed and cannot be implemented outside this crate.
204///
205/// # Examples
206///
207/// ```ignore
208/// use ndarray::Array2;
209/// use tenferro_ext_ndarray::try_ndarray_to_frontend;
210///
211/// let array = Array2::from_shape_vec((1, 2), vec![1.0_f64, 2.0]).unwrap();
212/// let tensor = try_ndarray_to_frontend(array).unwrap();
213/// assert_eq!(tensor.dims(), &[1, 2]);
214/// ```
215#[cfg(feature = "frontend")]
216pub trait FrontendScalar: Scalar + frontend::Sealed {
217    fn into_frontend_tensor(tensor: Tensor<Self>) -> tenferro::Tensor;
218}
219
220#[cfg(feature = "frontend")]
221impl FrontendScalar for f32 {
222    fn into_frontend_tensor(tensor: Tensor<Self>) -> tenferro::Tensor {
223        tenferro::Tensor::from_tensor(tensor)
224    }
225}
226
227#[cfg(feature = "frontend")]
228impl FrontendScalar for f64 {
229    fn into_frontend_tensor(tensor: Tensor<Self>) -> tenferro::Tensor {
230        tenferro::Tensor::from_tensor(tensor)
231    }
232}
233
234#[cfg(feature = "frontend")]
235impl FrontendScalar for num_complex::Complex32 {
236    fn into_frontend_tensor(tensor: Tensor<Self>) -> tenferro::Tensor {
237        tenferro::Tensor::from_tensor(tensor)
238    }
239}
240
241#[cfg(feature = "frontend")]
242impl FrontendScalar for num_complex::Complex64 {
243    fn into_frontend_tensor(tensor: Tensor<Self>) -> tenferro::Tensor {
244        tenferro::Tensor::from_tensor(tensor)
245    }
246}
247
248/// Fallibly converts an owned ndarray array into the public `tenferro::Tensor`
249/// frontend.
250///
251/// This helper is enabled by the `frontend` feature and is intentionally thin:
252/// it converts through the canonical typed `ndarray -> tenferro_tensor::Tensor`
253/// path first.
254///
255/// # Examples
256///
257/// ```ignore
258/// use ndarray::Array2;
259/// use tenferro_ext_ndarray::try_ndarray_to_frontend;
260///
261/// let array = Array2::from_shape_vec((1, 2), vec![1.0_f64, 2.0]).unwrap();
262/// let tensor = try_ndarray_to_frontend(array).unwrap();
263/// assert_eq!(tensor.dims(), &[1, 2]);
264/// ```
265#[cfg(feature = "frontend")]
266pub fn try_ndarray_to_frontend<T, S, D>(array: ArrayBase<S, D>) -> Result<tenferro::Tensor>
267where
268    T: FrontendScalar + Clone,
269    S: Data<Elem = T>,
270    D: Dimension,
271{
272    let tensor = try_ndarray_to_tensor(array)?;
273    Ok(T::into_frontend_tensor(tensor))
274}