briny_ai 0.3.0

A tiny & efficient AI inference engine
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

use crate::abstracted::internal::TensorLike;
use crate::manual::TensorFloat;
use crate::manual::tensors::WithGrad;

#[cfg(feature = "dyntensor")]
use crate::manual::tensors::{Tensor, TensorOps};

mod dense;
pub use dense::*;

#[cfg(feature = "dyntensor")]
fn reduce_rows_sum_to_bias(dz: &Tensor<TensorFloat>, out_len: usize) -> Tensor<TensorFloat> {
    let mut acc = alloc::vec![0.0; out_len]; // or a fixed array if OUT is const
    let dz_data = dz.data();
    let batch = dz_data.len() / out_len;

    for i in 0..batch {
        for j in 0..out_len {
            acc[j] += dz_data[i * out_len + j];
        }
    }

    Tensor::new(&[out_len], &acc) // however you normally construct 1D tensors
}

#[cfg(feature = "alloc")]
use alloc::boxed::Box;
#[cfg(not(feature = "alloc"))]
use box_closure::OpaqueFnOnce;

#[cfg(feature = "alloc")]
pub type BackFn<'a, In, Out> = Box<dyn FnOnce(Out) -> In + 'a>;

#[cfg(not(feature = "alloc"))]
pub type BackFn<'a, In, Out> = OpaqueFnOnce<'a, Out, In, box_closure::Align32<1024>>;

/// Internal utilities to operate on layers.
///
/// Obtained from `__LayoutMarker` at compile time when
/// using the `#[model]` proc-macro.
pub trait Layer<I, O> {
    type ParamA: TensorLike<TensorFloat>;
    type ParamB: TensorLike<TensorFloat>;

    fn run<'a>(&'a mut self, input: &'a WithGrad<I>) -> (O, BackFn<'a, I, O>);

    fn params(&mut self) -> (&mut WithGrad<Self::ParamA>, &mut WithGrad<Self::ParamB>);
}

/// Marker trait for layout dummies.
///
/// Converted to [`Layer<I, O>`] at compile time when
/// using the `#[model]` proc-macro.
/// The internal [`Layer`] trait
pub trait __LayoutMarker {}