rlx_tensor/interop.rs
1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! `ndarray` interop (feature `ndarray`).
6//!
7//! Migrate in: `Tensor::from(array)` for any `ndarray::Array<f32, D>` (any
8//! rank). Migrate out (with `eval`): `tensor.to_ndarray()` realizes the lazy
9//! graph and hands back an owned `ArrayD<f32>`.
10
11use ndarray::{Array, ArrayD, Dimension, IxDyn};
12
13use crate::Tensor;
14
15impl<D: Dimension> From<Array<f32, D>> for Tensor {
16 fn from(a: Array<f32, D>) -> Self {
17 Tensor::from(&a)
18 }
19}
20
21impl<D: Dimension> From<&Array<f32, D>> for Tensor {
22 fn from(a: &Array<f32, D>) -> Self {
23 let dims = a.shape().to_vec();
24 // `iter()` yields logical (row-major) order regardless of memory layout.
25 let data: Vec<f32> = a.iter().copied().collect();
26 Tensor::from_vec(data, dims)
27 }
28}
29
30#[cfg(feature = "eval")]
31impl Tensor {
32 /// Realize (compile + run) and return an owned `ndarray::ArrayD<f32>`.
33 pub fn to_ndarray(&self) -> ArrayD<f32> {
34 let dims = self.dims();
35 let data = self.to_vec();
36 ArrayD::from_shape_vec(IxDyn(&dims), data)
37 .expect("to_ndarray: realized data does not match shape")
38 }
39}