Skip to main content

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}