File size: 830 Bytes
712dbf0 |
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 |
# 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()
|