libmir-metal 0.1.0

Metal inference backend for libmir
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
56
57
58
59

use super::*;

#[test]
fn selects_largest_router_scores() -> Result<()> {
    let stream = Stream::new_gpu()?;
    let scores = (0_u16..128).map(f32::from).collect::<Vec<_>>();
    let scores = Array::from_f32(&scores, &[1, 1, 128])?;
    let scales = Array::from_f32(&vec![1.0; 128], &[128])?;
    let mut indices = scores
        .router_top_k(&scales, 8, &stream)?
        .indices
        .to_vec_u32_on_stream(&stream)?;
    indices.sort_unstable();
    assert_eq!(indices, (120_u32..128).collect::<Vec<_>>());
    Ok(())
}

#[test]
fn selects_largest_scores_for_each_router_token() -> Result<()> {
    let stream = Stream::new_gpu()?;
    let mut scores = (0_u16..128).map(f32::from).collect::<Vec<_>>();
    scores.extend((0_u16..128).rev().map(f32::from));
    let scores = Array::from_f32(&scores, &[1, 2, 128])?;
    let scales = Array::from_f32(&vec![1.0; 128], &[128])?;
    let indices = scores
        .router_top_k(&scales, 8, &stream)?
        .indices
        .to_vec_u32_on_stream(&stream)?;
    assert_eq!(indices.len(), 16);
    let mut first = indices[..8].to_vec();
    let mut second = indices[8..].to_vec();
    first.sort_unstable();
    second.sort_unstable();
    assert_eq!(first, (120_u32..128).collect::<Vec<_>>());
    assert_eq!(second, (0_u32..8).collect::<Vec<_>>());
    Ok(())
}

#[test]
fn unit_router_matches_an_explicit_all_ones_scale() -> Result<()> {
    let stream = Stream::new_gpu()?;
    let scores = (0_u16..128).map(|value| f32::from(value) / 17.0).collect::<Vec<_>>();
    let scores = Array::from_f32(&scores, &[1, 1, 128])?;
    let unit_factors = Array::from_f32(&vec![1.0; 128], &[128])?;
    let explicit_routing = scores.router_top_k(&unit_factors, 8, &stream)?;
    let unit = scores.router_top_k_unit(8, &stream)?;

    explicit_routing.indices.async_eval()?;
    explicit_routing.weights.async_eval()?;
    unit.indices.async_eval()?;
    unit.weights.async_eval()?;
    stream.synchronize()?;
    assert_eq!(
        explicit_routing.indices.to_vec_u32_on_stream(&stream)?,
        unit.indices.to_vec_u32_on_stream(&stream)?,
    );
    assert_eq!(explicit_routing.weights.to_vec_f32()?, unit.weights.to_vec_f32()?);
    Ok(())
}