vyre-libs 0.7.2

vyre Category A library ecosystem - pure-IR compositions over foundation IR and primitive-owned kernels
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
66
67
68
69
70
71
72
73
74
75

//! Parallel residual block: `out = x + attn_out + mlp_out`.
//!
//! Category A composition  -  residual stream addition.

use vyre_foundation::ir::{BufferAccess, BufferDecl, DataType, Expr, Program};

use crate::builder::build_indexed_map;

const OP_ID: &str = "vyre-libs::nn::parallel_residual_block";

/// Build parallel residual block (F32).
///
/// # Errors
/// Returns `Err` if n is zero.
pub fn parallel_residual_block(
    x: &str,
    attn_out: &str,
    mlp_out: &str,
    output: &str,
    n: u32,
) -> Result<Program, String> {
    if n == 0 {
        return Err("Fix: n=0".into());
    }
    let buffers = vec![
        BufferDecl::storage(x, 0, BufferAccess::ReadOnly, DataType::F32).with_count(n),
        BufferDecl::storage(attn_out, 1, BufferAccess::ReadOnly, DataType::F32).with_count(n),
        BufferDecl::storage(mlp_out, 2, BufferAccess::ReadOnly, DataType::F32).with_count(n),
        BufferDecl::output(output, 3, DataType::F32).with_count(n),
    ];
    Ok(build_indexed_map(
        OP_ID,
        buffers,
        output,
        n,
        [64, 1, 1],
        |i| {
            (
                i.clone(),
                Expr::add(
                    Expr::add(Expr::load(x, i.clone()), Expr::load(attn_out, i.clone())),
                    Expr::load(mlp_out, i),
                ),
            )
        },
    ))
}

inventory::submit! {
    vyre_foundation::operation::OperationRegistration {
        semantic_version: 1,
        signature: None,
        tier: vyre_foundation::operation::OperationTier::Library,
        laws: &[],
        tolerance: vyre_foundation::operation::TolerancePolicy::EXACT,
        id: OP_ID,
        build: Some(|| {
            parallel_residual_block("x", "attn", "mlp", "out", 4)
                .unwrap_or_else(|error| crate::invalid_program(OP_ID, format!("Fix: parallel_residual_block fixture must build: {error}")))
        }),
        test_inputs: Some(|| {
            let f = vyre_primitives::wire::pack_f32_slice;
            vec![vec![
                f(&[1.0, 2.0, 3.0, 4.0]), f(&[0.1, 0.2, 0.3, 0.4]),
                f(&[0.01, 0.02, 0.03, 0.04]),
            ]]
        }),
        expected_output: Some(|| {
            let out = [1.11_f32, 2.22, 3.33, 4.44];
            let bytes = vyre_primitives::wire::pack_f32_slice(&out);
            vec![vec![bytes]]
        }),
        category: Some("nn"),
    }
}