Merge pull request #76 from adtzlr/add-dual2real

Add `Tensor.dual2real(like=None)`

src/tensortrax/_tensor.py

@@ -309,6 +309,9 @@ def reshape(self, *shape, order="C"):
     def squeeze(self, axis=None):
         return squeeze(self, axis=axis)
 
+    def dual2real(self, like):
+        return dual2real(self, like=like)
+
     __radd__ = __add__
     __rmul__ = __mul__
     __array_ufunc__ = None
@@ -331,6 +334,16 @@ def broadcast_to(A, shape):
         return _broadcast_to(A)
 
 
+def dual2real(A, like=None):
+    """Return a new Tensor with old-dual data as new-real values,
+    with `ntrax` derived by `like`."""
+
+    ndual = like.ndual - len(like.shape)
+    ntrax = A.ntrax - ndual
+
+    return Tensor(x=A.δx, δx=A.Δδx, Δx=A.Δδx, ndual=min(ndual, ntrax), ntrax=ntrax)
+
+
 def ravel(A, order="C"):
     if isinstance(A, Tensor):
         δtrax = δ(A).shape[len(A.shape) :]

src/tensortrax/math/__init__.py

@@ -19,6 +19,7 @@
     cosh,
     diagonal,
     dot,
+    dual2real,
     einsum,
     exp,
     hstack,

src/tensortrax/math/_math_tensor.py

@@ -15,6 +15,7 @@
     Δ,
     Δδ,
     broadcast_to,
+    dual2real,
     einsum,
     f,
     matmul,

tests/test_dual2real.py

@@ -0,0 +1,36 @@
+import numpy as np
+
+import tensortrax as tr
+import tensortrax.math as tm
+
+
+def test_dual2real():
+    np.random.seed(34563)
+    x = (np.random.rand(3, 3) - 0.5) / 10 + np.eye(3)
+
+    # init a Tensor with `hessian=True`
+    F = tr.Tensor(x)
+    F.init(hessian=True)
+
+    # perform some math operations
+    C = F.T() @ F
+    J = tm.linalg.det(F)
+    W = tm.trace(J**(-2 / 3) * C) - 3
+    eta = 1 - 1 / 3 * tm.tanh(W / 8)
+
+    # set old dual data as new real values (i.e. obtain the gradient)
+    P = W.dual2real(like=F)
+    tm.dual2real(W, like=F)
+
+    # perform some more math with a derived Tensor involved
+    Q = eta * P
+
+    # take the gradient
+    A = tr.δ(Q)
+
+    assert P.shape == (3, 3)
+    assert Q.shape == (3, 3)
+    assert A.shape == (3, 3, 3, 3)
+
+if __name__ == "__main__":
+    test_dual2real()