mlex 0.1.3

Safe, idiomatic Rust runtime for Apple MLX: quantized LLM inference (Qwen, Gemma4, NemotronH, DharaAR, ...) with multi-modal support
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

# Copyright © 2023 Apple Inc.

import numpy as np
from time_utils import time_fn


def time_add():
    a = np.ones((100, 100, 10), dtype=np.float32)
    b = np.ones((100, 100, 10), dtype=np.float32)
    time_fn(np.add, a, b)


def time_matmul():
    a = np.random.rand(1000, 500).astype(np.float32)
    b = np.random.rand(500, 1000).astype(np.float32)
    time_fn(np.matmul, a, b)


def time_exp():
    a = np.random.randn(1000, 100).astype(np.float32)
    time_fn(np.exp, a)


def time_take():
    a = np.random.rand(10000, 500)
    ids = np.random.randint(0, 10000, (20, 10))
    ids = [idx.reshape(-1) for idx in np.split(ids, 20)]

    def random_take():
        return [np.take(a, idx, 0) for idx in ids]

    time_fn(random_take)


if __name__ == "__main__":
    time_add()
    time_matmul()
    time_exp()
    time_take()