| 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 | |