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

//! Element-wise residual-stream addition.

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

use super::unary::typed_binary_activation_program;

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

/// Compute `output[i] = residual[i] + branch[i]` in F32.
#[must_use]
pub fn residual_add(residual: &str, branch: &str, output: &str, n: u32) -> Program {
    build_residual_add(residual, branch, output, n, DataType::F32)
}

/// Add typed residual and branch tensors with F32 arithmetic.
///
/// # Errors
///
/// Returns `Err` for an empty vector or a non-floating activation dtype.
pub fn residual_add_typed(
    residual: &str,
    branch: &str,
    output: &str,
    n: u32,
    dtype: DataType,
) -> Result<Program, String> {
    if n == 0 {
        return Err("Fix: residual_add_typed requires n > 0".to_string());
    }
    if !matches!(dtype, DataType::F16 | DataType::BF16 | DataType::F32) {
        return Err(format!(
            "Fix: residual_add_typed supports F16, BF16, or F32 tensors; got {dtype:?}"
        ));
    }
    Ok(build_residual_add(residual, branch, output, n, dtype))
}

fn build_residual_add(
    residual: &str,
    branch: &str,
    output: &str,
    n: u32,
    dtype: DataType,
) -> Program {
    if n == 0 {
        return crate::invalid_program(OP_ID, "Fix: residual_add requires n > 0");
    }
    typed_binary_activation_program(OP_ID, residual, branch, output, n, dtype, Expr::add)
}
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(|| residual_add("residual", "branch", "output", 4)),
        test_inputs: Some(|| {
            vec![vec![
                vyre_primitives::wire::pack_f32_slice(&[1.0, -2.0, 3.5, 0.0]),
                vyre_primitives::wire::pack_f32_slice(&[0.5, 4.0, -1.5, -0.0]),
            ]]
        }),
        expected_output: Some(|| {
            vec![vec![vyre_primitives::wire::pack_f32_slice(&[
                1.5, 2.0, 2.0, 0.0,
            ])]]
        }),
        category: Some("nn"),
    }
}