added uexp

src/tricycle_v2/unary.py

@@ -127,3 +127,17 @@ def umin(tensor: Tensor, constant: float) -> Tensor:
     result.args = (tensor,)
     result.back_fn = (partial(einsum, subscripts, indicator),)
     return result
+
+
+def uexp(tensor: Tensor) -> Tensor:
+    """
+    Raise every element of a tensor to the power of e
+    """
+    result = to_tensor(np.exp(tensor))
+
+    indices = ascii_letters[: len(tensor.shape)]
+    subscripts = f"{indices},{indices}->{indices}"
+
+    result.back_fn = (partial(einsum, subscripts, result),)
+    result.args = (tensor,)
+    return result

tests/test_unary_ops.py

@@ -1,7 +1,7 @@
 import numpy as np
 
 from tricycle_v2.ops import einsum, to_tensor
-from tricycle_v2.unary import uadd, udiv, umax, umul, upow, usub, umin
+from tricycle_v2.unary import uadd, udiv, uexp, umax, umin, umul, upow, usub
 
 
 def test_can_add():
@@ -152,14 +152,42 @@ def test_can_umin():
     in_tensor = to_tensor(np.arange(12).reshape(3, 4))
     out_tensor = umin(in_tensor, 4)
 
+    assert out_tensor.shape == (3, 4)
+    assert np.allclose(out_tensor, np.array([[0, 1, 2, 3], [4, 4, 4, 4], [4, 4, 4, 4]]))
+
+    out_tensor.backward()
+
+    assert np.allclose(
+        in_tensor.grad,
+        np.array([[1.0, 1.0, 1.0, 1.0], [1.0, 0, 0, 0], [0, 0, 0, 0]]),
+    )
+
+
+def test_can_uexp():
+    in_tensor = to_tensor(np.arange(12).reshape(3, 4))
+    out_tensor = uexp(in_tensor)
+
     assert out_tensor.shape == (3, 4)
     assert np.allclose(
-        out_tensor, np.array([[0, 1, 2, 3], [4, 4, 4, 4], [4, 4, 4, 4]])
+        out_tensor,
+        np.array(
+            [
+                [1.0, np.exp(1), np.exp(2), np.exp(3)],
+                [np.exp(4), np.exp(5), np.exp(6), np.exp(7)],
+                [np.exp(8), np.exp(9), np.exp(10), np.exp(11)],
+            ]
+        ),
     )
 
     out_tensor.backward()
 
     assert np.allclose(
         in_tensor.grad,
-        np.array([[1.0, 1.0, 1.0, 1.0], [1.0, 0, 0, 0], [0, 0, 0, 0]]),
+        np.array(
+            [
+                [np.exp(0), np.exp(1), np.exp(2), np.exp(3)],
+                [np.exp(4), np.exp(5), np.exp(6), np.exp(7)],
+                [np.exp(8), np.exp(9), np.exp(10), np.exp(11)],
+            ]
+        ),
     )