github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-09 15:08:02 -05:00
Code Issues Packages Projects Releases Wiki Activity

[Bounty] Vectorize Transcendental (#9058)

Browse Source 
* init

* cast everythig right

* more casting

* install pillow in test

* quick tests

* simplify

* quick tests

* delete test

* tests

* fix import error

* add vec to ldexp3k

* vec for bitcast

* some helper tests

* high level tests

* clean tests

* change tolerance so cuda passes

* ruff passes

* remove tests for transcendental helpers

* ruff passes

* make exponent in power vectorized

* fix pow test

* add newline

* add vec dtype to ilogb2k

* comment + clean up

* ruff

---------

Co-authored-by: chenyu <chenyu@fastmail.com>
Co-authored-by: George Hotz <72895+geohot@users.noreply.github.com>
This commit is contained in:
Eitan Turok
2025-02-28 09:47:25 +02:00
committed by GitHub
parent 8ae215dd3d
commit d657d5f754
4 changed files with 54 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
test/test_transcendental.py
View File

@@ -128,5 +128,33 @@ class TestTranscendentalSchedule(unittest.TestCase):
c = c.exp2()
check_schedule(c, 1)
class TestTranscendentalVectorized(unittest.TestCase):
def _vectorized_data(self, low, high, vec_size):
np_data = np.linspace(low, high, num=(128 // vec_size) * vec_size, dtype=np.float32).reshape(-1, vec_size)
data = Tensor(np_data, dtype=dtypes.float32.vec(vec_size))
return data, np_data
def _test_vectorized_op(self, fxn, np_fxn, data_range, vec_size, param_range=None):
data, np_data = self._vectorized_data(data_range[0], data_range[1], vec_size)
if param_range:
param, np_param = self._vectorized_data(param_range[0], param_range[1], vec_size)
out, np_out = fxn(data, param), np_fxn(np_data, np_param)
else:
out, np_out = fxn(data), np_fxn(np_data)
np.testing.assert_allclose(out.numpy(), np_out, rtol=1e-4)
def test_exp2_vectorized(self):
for vec_size in [1,2,3,4,5,127,128]: self._test_vectorized_op(Tensor.exp2, np.exp2, (-100, 100), vec_size)
def test_log2_vectorized(self):
for vec_size in [1,2,3,4,5,127,128]: self._test_vectorized_op(Tensor.log2, np.log2, (0.001, 200), vec_size)
def test_sin_vectorized(self):
for vec_size in [1,2,3,4,5,127,128]: self._test_vectorized_op(Tensor.sin, np.sin, (-100, 100), vec_size)
def test_pow_vectorized(self):
# np.pow returns nan for negative values raised to a non-integral power
for vec_size in [1,2,3,4,5,127,128]: self._test_vectorized_op(Tensor.pow, np.pow, (0.001, 200), vec_size, param_range=(-10, 10))
if __name__ == '__main__':
unittest.main()