Skip to main content

neura_decoder/
parameter.rs

1use neura_abi::Element;
2use neura_gguf::{Archive, Tensor};
3use neura_graph::{Graph, Init, Shape, Value};
4
5pub fn matrix(graph: &Graph, source: &impl Archive, name: &str, rows: u32, columns: u32) -> Value {
6    let held = held(source, name);
7    assert_eq!(
8        (held.rank, held.shape[0], held.shape[1]),
9        (2, u64::from(columns), u64::from(rows)),
10        "the tensor {name} of this gguf walks {} dimensions of {:?} where {rows} rows of {columns} numbers walk [1, 1, {rows}, {columns}]",
11        held.rank,
12        held.shape,
13    );
14    declare(graph, name, Shape::matrix(rows, columns), held.element)
15}
16
17pub fn vector(graph: &Graph, source: &impl Archive, name: &str, width: u32) -> Value {
18    let held = held(source, name);
19    assert_eq!(
20        (held.rank, held.shape[0]),
21        (1, u64::from(width)),
22        "the tensor {name} of this gguf walks {} dimensions of {:?} where {width} numbers walk [1, 1, 1, {width}]",
23        held.rank,
24        held.shape,
25    );
26    declare(graph, name, Shape::vector(width), held.element)
27}
28
29fn held(source: &impl Archive, name: &str) -> Tensor {
30    source
31        .tensor(name)
32        .unwrap_or_else(|| panic!("this gguf holds no tensor named {name}"))
33}
34
35fn declare(graph: &Graph, name: &str, shape: Shape, element: Element) -> Value {
36    if element.block().is_some() {
37        graph.named_block_quantized_parameter(name, shape, Init::Zero, element)
38    } else {
39        graph.named_parameter(name, shape, Init::Zero, element)
40    }
41}