clean up transcendental frexp [pr] (#7384)

also added some unit tests for frexp
This commit is contained in:
chenyu
2024-10-29 18:51:37 -04:00
committed by GitHub
parent 99b82f5708
commit 6bf38c35e5
3 changed files with 34 additions and 12 deletions

View File

@@ -3,7 +3,7 @@ import numpy as np
from tinygrad import dtypes
from tinygrad.ops import UOp, UOps
from tinygrad.codegen.uopgraph import full_graph_rewrite
from tinygrad.codegen.transcendental import payne_hanek_reduction, cody_waite_reduction
from tinygrad.codegen.transcendental import payne_hanek_reduction, cody_waite_reduction, frexp
from tinygrad.codegen.linearize import linearize_uop
from tinygrad.runtime.ops_python import PythonProgram, PythonRenderer, PythonCompiler, PythonAllocator
@@ -28,5 +28,26 @@ class TestReduction(unittest.TestCase):
# TODO: should q be in [0, 1, 2, 3]?
np.testing.assert_equal(q, 12)
def test_frexp(self):
mantissa, exponent = (self._run_uop(u) for u in frexp(UOp.const(dtypes.float64, 0.0)))
np.testing.assert_equal(mantissa, 0.0)
np.testing.assert_equal(exponent, 0)
mantissa, exponent = (self._run_uop(u) for u in frexp(UOp.const(dtypes.float64, 1.0)))
np.testing.assert_equal(mantissa, 0.5)
np.testing.assert_equal(exponent, 1)
mantissa, exponent = (self._run_uop(u) for u in frexp(UOp.const(dtypes.float64, -1.0)))
np.testing.assert_equal(mantissa, 0.5)
np.testing.assert_equal(exponent, 1)
mantissa, exponent = (self._run_uop(u) for u in frexp(UOp.const(dtypes.float64, 2.0)))
np.testing.assert_equal(mantissa, 0.5)
np.testing.assert_equal(exponent, 2)
mantissa, exponent = (self._run_uop(u) for u in frexp(UOp.const(dtypes.float64, 5.0)))
np.testing.assert_equal(mantissa, 0.625)
np.testing.assert_equal(exponent, 3)
if __name__ == '__main__':
unittest.main()