rlx-gemma 0.2.11

Gemma / Gemma 2 causal LMs for RLX
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

// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, version 3.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program. If not, see <https://www.gnu.org/licenses/>.

//! Backend capability checks for Gemma.2.

use crate::config::GemmaConfig;
use anyhow::Result;
use rlx_core::validate_lm_device;
pub use rlx_core::{LM_DEVICE_NAMES, STANDARD_DEVICE_NAMES, STANDARD_DEVICES, pick_lm_device};
use rlx_runtime::Device;

pub fn validate_device(cfg: &GemmaConfig, device: Device, packed_weights: bool) -> Result<()> {
    let _ = (cfg, packed_weights);
    validate_lm_device("gemma", device)
}

#[cfg(test)]
mod tests {
    use super::*;
    use rlx_core::STANDARD_DEVICES;

    fn tiny_cfg() -> GemmaConfig {
        GemmaConfig::tiny_test()
    }

    #[test]
    fn all_standard_backends_allowed() {
        let cfg = tiny_cfg();
        for dev in STANDARD_DEVICES {
            validate_device(&cfg, *dev, false).unwrap();
            validate_device(&cfg, *dev, true).unwrap();
        }
    }

    #[cfg(feature = "coreml")]
    #[test]
    fn coreml_ane_allowed() {
        let cfg = tiny_cfg();
        validate_device(&cfg, Device::Ane, true).unwrap();
    }
}