File size: 968 Bytes
ae9e4fe | 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 | import numpy as np
from mla import Tensor, gradcheck, exp, log, rsqrt, silu, gelu, softmax
def _rand(*shape):
rng = np.random.default_rng(1)
return Tensor(rng.standard_normal(shape))
def _pos(*shape):
rng = np.random.default_rng(2)
return Tensor(np.abs(rng.standard_normal(shape)) + 0.5)
def test_exp():
ok, err = gradcheck(lambda a: exp(a), [_rand(3, 4)])
assert ok, err
def test_log():
ok, err = gradcheck(lambda a: log(a), [_pos(3, 4)])
assert ok, err
def test_rsqrt():
ok, err = gradcheck(lambda a: rsqrt(a), [_pos(3, 4)])
assert ok, err
def test_silu():
ok, err = gradcheck(lambda a: silu(a), [_rand(3, 4)])
assert ok, err
def test_gelu():
ok, err = gradcheck(lambda a: gelu(a), [_rand(3, 4)])
assert ok, err
def test_softmax():
w = Tensor(np.random.default_rng(3).standard_normal((3, 5)))
ok, err = gradcheck(lambda a: softmax(a, axis=-1).mul(w), [_rand(3, 5)])
assert ok, err
|