rlx-tensor 0.2.18

Native ndarray alternative for RLX — NumPy-style tensors that trace, fuse, and run on any backend via the rlx compiler
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39

// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
// SPDX-License-Identifier: MIT OR Apache-2.0

//! `ndarray` interop (feature `ndarray`).
//!
//! Migrate in: `Tensor::from(array)` for any `ndarray::Array<f32, D>` (any
//! rank). Migrate out (with `eval`): `tensor.to_ndarray()` realizes the lazy
//! graph and hands back an owned `ArrayD<f32>`.

use ndarray::{Array, ArrayD, Dimension, IxDyn};

use crate::Tensor;

impl<D: Dimension> From<Array<f32, D>> for Tensor {
    fn from(a: Array<f32, D>) -> Self {
        Tensor::from(&a)
    }
}

impl<D: Dimension> From<&Array<f32, D>> for Tensor {
    fn from(a: &Array<f32, D>) -> Self {
        let dims = a.shape().to_vec();
        // `iter()` yields logical (row-major) order regardless of memory layout.
        let data: Vec<f32> = a.iter().copied().collect();
        Tensor::from_vec(data, dims)
    }
}

#[cfg(feature = "eval")]
impl Tensor {
    /// Realize (compile + run) and return an owned `ndarray::ArrayD<f32>`.
    pub fn to_ndarray(&self) -> ArrayD<f32> {
        let dims = self.dims();
        let data = self.to_vec();
        ArrayD::from_shape_vec(IxDyn(&dims), data)
            .expect("to_ndarray: realized data does not match shape")
    }
}