1use crate::write::{PackedWeight, WriteOptions, write_package};
8use anyhow::Result;
9use rlx_ir::Graph;
10use std::path::Path;
11
12#[derive(Debug, Clone)]
14pub struct BakeWeight {
15 pub name: String,
16 pub shape: Vec<usize>,
17 pub encoding: String,
19 pub data: Vec<u8>,
20}
21
22impl BakeWeight {
23 fn layout_for_encoding(encoding: &str) -> &'static str {
24 if encoding.starts_with("gguf_") {
25 "bt_nk"
26 } else {
27 "row_major"
28 }
29 }
30}
31
32pub fn package_from_bake(
34 out: impl AsRef<Path>,
35 graph: &Graph,
36 weights: &[BakeWeight],
37 mut opts: WriteOptions,
38) -> Result<()> {
39 if opts.name == "model" {
40 if !graph.name.is_empty() {
41 opts.name = graph.name.clone();
42 }
43 }
44 if opts.producer.is_none() {
45 opts.producer = Some("rlx-bake".into());
46 }
47 let packed: Vec<PackedWeight> = weights
48 .iter()
49 .map(|w| PackedWeight {
50 name: w.name.clone(),
51 shape: w.shape.clone(),
52 scheme: w.encoding.clone(),
53 layout: BakeWeight::layout_for_encoding(&w.encoding).into(),
54 data: w.data.clone(),
55 rank: None,
56 tier: crate::tier::StorageTier::Hot,
57 })
58 .collect();
59 write_package(out, graph, &packed, &opts)
60}