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
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89

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

//! Python-scalar promotion rules mirrored from `pyrlx/dsl.py`.

use rlx_ir::{DType, Graph, GraphExt, NodeId};

/// Maximum integer magnitude representable exactly as f64.
pub(crate) const MAX_EXACT_INT: i128 = 1 << 53;

pub(crate) fn promote_scalar(g: &mut Graph, value: Scalar, dtype_hint: DType) -> NodeId {
    match value {
        Scalar::Bool(b) => g.constant(if b { 1.0 } else { 0.0 }, DType::Bool),
        Scalar::Int(v) => {
            let fv = int_to_f64(v).expect("integer literal exceeds exact f64 range");
            if dtype_hint.is_integral() {
                g.constant(fv, dtype_hint)
            } else {
                g.constant(fv, DType::F32)
            }
        }
        Scalar::Float(v) => {
            if dtype_hint.is_integral() {
                g.constant(v, dtype_hint)
            } else {
                g.constant(v, DType::F32)
            }
        }
    }
}

fn int_to_f64(v: i64) -> Result<f64, ()> {
    if v.unsigned_abs() > MAX_EXACT_INT as u64 {
        return Err(());
    }
    Ok(v as f64)
}

/// A literal operand in elementwise binary ops.
#[derive(Debug, Clone, Copy)]
pub(crate) enum Scalar {
    Bool(bool),
    Int(i64),
    Float(f64),
}

impl From<bool> for Scalar {
    fn from(v: bool) -> Self {
        Self::Bool(v)
    }
}

impl From<i32> for Scalar {
    fn from(v: i32) -> Self {
        Self::Int(v as i64)
    }
}

impl From<i64> for Scalar {
    fn from(v: i64) -> Self {
        Self::Int(v)
    }
}

impl From<f32> for Scalar {
    fn from(v: f32) -> Self {
        Self::Float(v as f64)
    }
}

impl From<f64> for Scalar {
    fn from(v: f64) -> Self {
        Self::Float(v)
    }
}

trait IntegralDType {
    fn is_integral(&self) -> bool;
}

impl IntegralDType for DType {
    fn is_integral(&self) -> bool {
        matches!(
            self,
            DType::I8 | DType::I16 | DType::I32 | DType::I64 | DType::U8 | DType::U32
        )
    }
}