Skip to main content

rlx_pkg/
from_bake.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! Build an `.rlxp` from bake-style graph + weight table (no bake crate dep).
6
7use crate::write::{PackedWeight, WriteOptions, write_package};
8use anyhow::Result;
9use rlx_ir::Graph;
10use std::path::Path;
11
12/// Minimal weight row matching bake’s table (avoids depending on `rlx-bake`).
13#[derive(Debug, Clone)]
14pub struct BakeWeight {
15    pub name: String,
16    pub shape: Vec<usize>,
17    /// Bake encoding string: `f32`, `gguf_tq2_0`, `gguf_q8_0`, …
18    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
32/// Write a package from a baked graph and weight table.
33pub 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}