add exp2 (#2192)

examples/f16_w_uint32.py

@@ -0,0 +1,42 @@
+import numpy as np
+from tinygrad.tensor import Tensor
+from tinygrad.helpers import dtypes
+from tinygrad.ops import Device
+
+# TODO: will be better when tinygrad does math in the target dtype, can remove the floor and use a mul
+def bit_extract(x, s, e) -> Tensor:
+  # extract the top bits we don't want
+  top_bits = (x / (1<<(s+1))).floor() * (1<<(s+1))
+  x = (x - top_bits) / (1<<e)
+  return x.contiguous()
+
+def u16_to_f16(x):
+  sign = bit_extract(x, 15, 15).float()
+  exponent = bit_extract(x, 14, 10).float()
+  fraction = bit_extract(x, 9, 0).float()
+  return sign.where(-1, 1) * exponent.where((exponent - 15).exp2() * (1 + fraction / 0x400), 6.103515625e-5 * (fraction / 0x400))
+
+def u32_to_f16(oo):
+  oo1 = (oo/0x10000).floor().contiguous()
+  # TODO: this is wrong and unextractable until we do this math in u32
+  oo2 = (oo-(oo1*0x10000)).floor().contiguous()
+  f1 = u16_to_f16(oo1)
+  f2 = u16_to_f16(oo2)
+  return Tensor.cat(f2.reshape(-1, 1), f1.reshape(-1, 1), dim=1).flatten()
+
+if __name__ == "__main__":
+  # random float16
+  Tensor.manual_seed(2)
+  a = Tensor.randn(100, dtype=dtypes.float16)
+
+  # this converts it to u32 on disk
+  oo = a.to("disk:/tmp/f16").cast(dtypes.uint32)[:50].to(Device.DEFAULT).realize()
+
+  # convert to 2xf16 using tinygrad math ops
+  f16 = u32_to_f16(oo)
+
+  ref = a.numpy()
+  out = f16.numpy().astype(np.float16)
+  print(ref-out)
+
+  np.testing.assert_allclose(ref, out)

test/test_ops.py

@@ -340,6 +340,9 @@ class TestOps(unittest.TestCase):
   def test_exp(self):
     helper_test_op([(45,65)], lambda x: torch.exp(x), Tensor.exp)
     helper_test_op([()], lambda x: torch.exp(x), Tensor.exp)
+  def test_exp2(self):
+    helper_test_op([(45,65)], lambda x: torch.exp2(x), Tensor.exp2)
+    helper_test_op([()], lambda x: torch.exp2(x), Tensor.exp2)
   def test_sign(self):
     helper_test_op([(45,65)], lambda x: torch.sign(x), Tensor.sign)
     helper_test_op([()], lambda x: torch.sign(x), Tensor.sign)

tinygrad/tensor.py

@@ -574,6 +574,7 @@ class Tensor:
   def log(self): return mlops.Log.apply(self)
   def log2(self): return mlops.Log.apply(self)/math.log(2)
   def exp(self): return mlops.Exp.apply(self)
+  def exp2(self): return mlops.Exp.apply(self*math.log(2))
   def relu(self): return mlops.Relu.apply(self)
   def sigmoid(self): return mlops.Sigmoid.apply(self)
   def sin(self): return mlops.Sin.apply(self)