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

//! Element-wise sigmoid output gate.

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

use super::unary::typed_sigmoid_gate_program;

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

/// Multiply `branch` by `sigmoid(gate_logits)` element-wise in F32.
#[must_use]
pub fn sigmoid_gate(gate_logits: &str, branch: &str, output: &str, n: u32) -> Program {
    build_sigmoid_gate(gate_logits, branch, output, n, DataType::F32)
}

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

fn build_sigmoid_gate(
    gate_logits: &str,
    branch: &str,
    output: &str,
    n: u32,
    dtype: DataType,
) -> Program {
    if n == 0 {
        return crate::invalid_program(OP_ID, "Fix: sigmoid_gate requires n > 0");
    }
    typed_sigmoid_gate_program(OP_ID, gate_logits, branch, output, n, dtype, false)
}
inventory::submit! {
    vyre_foundation::operation::OperationRegistration {
        semantic_version: 1,
        signature: None,
        tier: vyre_foundation::operation::OperationTier::Library,
        laws: &[],
        tolerance: vyre_foundation::operation::TolerancePolicy::f32_ulp(2),
        id: OP_ID,
        build: Some(|| sigmoid_gate("gate", "branch", "output", 4)),
        test_inputs: Some(|| {
            vec![vec![
                vyre_primitives::wire::pack_f32_slice(&[0.0, 1.0, -1.0, 100.0]),
                vyre_primitives::wire::pack_f32_slice(&[8.0, 2.0, -2.0, -7.0]),
            ]]
        }),
        expected_output: Some(|| {
            let gate = [0.0_f32, 1.0, -1.0, 100.0];
            let branch = [8.0_f32, 2.0, -2.0, -7.0];
            let output = std::array::from_fn::<_, 4, _>(|index| {
                branch[index] / (1.0 + (-gate[index]).exp())
            });
            vec![vec![vyre_primitives::wire::pack_f32_slice(&output)]]
        }),
        category: Some("nn"),
    }
}