Fix Tensor.init(gradient=True, δx=δx) (#123)

* Alternative `init` method of Tensor for gvp's

* Update _tensor.py

* Update _tensor.py

Author: adtzlr

src/tensortrax/_tensor.py

@@ -255,7 +255,11 @@ def init(self, gradient=False, hessian=False, sym=False, δx=None, Δx=None):
                     δx = np.zeros(self.size**2).reshape(shape)
                 else:
                     δx = np.eye(self.size).reshape(shape)
-                Δx = δx.copy()
+
+            else:
+                δx = δx.reshape(*self.shape, *self.trax)
+
+            Δx = δx.copy()
 
         elif hessian:
             # add additional trailing axes for dual values
@@ -454,7 +458,7 @@ def ravel(self, order="C"):
 
     def reshape(self, *shape, order="C"):
         "Gives a new shape to an array without changing its data."
-        return reshape(self, newshape=shape, order=order)
+        return reshape(self, shape, order=order)
 
     def squeeze(self, axis=None):
         "Remove axes of length one."
@@ -572,21 +576,21 @@ def squeeze(A, axis=None):
         return np.squeeze(A, axis=axis)
 
 
-def reshape(A, newshape, order="C"):
+def reshape(A, shape, order="C"):
     "Gives a new shape to an array without changing its data."
     if isinstance(A, Tensor):
         δtrax = δ(A).shape[len(A.shape) :]
         Δtrax = Δ(A).shape[len(A.shape) :]
         Δδtrax = Δδ(A).shape[len(A.shape) :]
         return Tensor(
-            x=f(A).reshape(*newshape, *A.trax, order=order),
-            δx=δ(A).reshape(*newshape, *δtrax, order=order),
-            Δx=Δ(A).reshape(*newshape, *Δtrax, order=order),
-            Δδx=Δδ(A).reshape(*newshape, *Δδtrax, order=order),
+            x=f(A).reshape(*shape, *A.trax, order=order),
+            δx=δ(A).reshape(*shape, *δtrax, order=order),
+            Δx=Δ(A).reshape(*shape, *Δtrax, order=order),
+            Δδx=Δδ(A).reshape(*shape, *Δδtrax, order=order),
             ntrax=A.ntrax,
         )
     else:
-        return np.reshape(A, newshape=newshape, order=order)
+        return np.reshape(A, shape, order=order)
 
 
 def einsum4(subscripts, *operands):

tests/test_hessian_vector_product.py

@@ -50,5 +50,17 @@ def test_hvp():
             assert not np.any(np.isnan(hvp))
 
 
+def test_gvp():
+    F = np.diag([2, 1.15, 1.15])
+    F[:, 0] += np.array([0, -0.15, -0.15])
+    δF = np.vstack([np.zeros(3), np.zeros(3), np.array([-0.04, 0.04, 0.04])])
+
+    δψ = tr.gradient_vector_product(neo_hooke)(F, δx=δF)
+    δψ_reference = tm.special.ddot(tr.gradient(neo_hooke)(F), δF)
+
+    assert np.isclose(δψ, δψ_reference)
+
+
 if __name__ == "__main__":
     test_hvp()
+    test_gvp()