vyre-libs 0.6.3

vyre Category A library ecosystem - pure-IR compositions over vyre-ops hardware primitives
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
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65

//! MuonEq-R: Row-normalized Muon optimizer (F32).
//!
//! Muon + `scale = max(1, rows/cols)^0.5` row normalization.

use vyre::ir::Program;

use crate::nn::optim::muon_core::muon_step_program;

const OP_ID: &str = "vyre-libs::optim::muoneq_r";

/// MuonEq-R step (F32).
///
/// Same as `muon_update` but with row-norm scaling baked in.
#[must_use]
pub fn muoneq_r(
    params: &str,
    grads: &str,
    momentum_buf: &str,
    output: &str,
    n: u32,
    rows: u32,
    cols: u32,
    lr: f32,
    momentum: f32,
) -> Program {
    // scale = max(1, rows/cols)^0.5
    let ratio = (rows as f32) / (cols as f32);
    let scale = ratio.max(1.0).sqrt();
    muon_step_program(
        OP_ID,
        params,
        grads,
        momentum_buf,
        output,
        n,
        scale * lr,
        momentum,
    )
}

inventory::submit! {
    crate::harness::OpEntry {
        id: OP_ID,
        build: || muoneq_r("params", "grads", "momentum", "output", 4, 4, 2, 0.02, 0.95),
        test_inputs: Some(|| {
            let to_f32 = |w: &[f32]| vyre_primitives::wire::pack_f32_slice(w);
            vec![vec![
                to_f32(&[1.0, 2.0, 3.0, 4.0]),
                to_f32(&[0.1, 0.2, 0.3, 0.4]),
                to_f32(&[0.0, 0.0, 0.0, 0.0]),
            ]]
        }),
        expected_output: Some(|| {
            vec![vec![
                vec![
                    205, 204, 204, 61, 205, 204, 76, 62, 154, 153, 153, 62, 205, 204, 204, 62,
                ],
                vec![
                    64, 239, 125, 63, 64, 239, 253, 63, 112, 115, 62, 64, 64, 239, 125, 64,
                ],
            ]]
        }),
        category: Some("nn"),
    }
}