rlx-runtime 0.2.14

RLX runtime — feature-gated backends, session API, compile+execute pipeline
Documentation 
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

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

//! `Op::Sort` / `Op::ArgSort` on CPU — values + indices, ascending/descending,
//! along either axis.

use rlx_ir::{DType, Graph, Shape};
use rlx_runtime::{Device, Session};

fn run(x: &[f32], dims: &[usize], axis: usize, descending: bool, arg: bool) -> Vec<f32> {
    let mut g = Graph::new("sort");
    let xi = g.input("x", Shape::new(dims, DType::F32));
    let out = if arg {
        g.argsort(xi, axis, descending, Shape::new(dims, DType::F32))
    } else {
        g.sort(xi, axis, descending, Shape::new(dims, DType::F32))
    };
    g.set_outputs(vec![out]);
    Session::new(Device::Cpu)
        .compile(g)
        .run(&[("x", x)])
        .pop()
        .unwrap()
}

#[test]
fn sort_argsort_cpu() {
    // [[3, 1, 2], [0.5, 4, -1]]
    let x = vec![3.0, 1.0, 2.0, 0.5, 4.0, -1.0];

    assert_eq!(
        run(&x, &[2, 3], 1, false, false),
        vec![1.0, 2.0, 3.0, -1.0, 0.5, 4.0]
    );
    assert_eq!(
        run(&x, &[2, 3], 1, false, true),
        vec![1.0, 2.0, 0.0, 2.0, 0.0, 1.0]
    );
    assert_eq!(
        run(&x, &[2, 3], 1, true, false),
        vec![3.0, 2.0, 1.0, 4.0, 0.5, -1.0]
    );
    // Axis 0: cols [3,0.5]→[0.5,3], [1,4]→[1,4], [2,-1]→[-1,2].
    assert_eq!(
        run(&x, &[2, 3], 0, false, false),
        vec![0.5, 1.0, -1.0, 3.0, 4.0, 2.0]
    );
}