Merge pull request #22 from adtzlr/add-reshape

Add reshape

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "tensortrax"
-version = "0.2.1"
+version = "0.2.2"
 description = "Math on (Hyper-Dual) Tensors with Trailing Axes"
 readme = "README.md"
 requires-python = ">=3.7"

setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = tensortrax
-version = 0.2.1
+version = 0.2.2
 author = Andreas Dutzler
 author_email = a.dutzler@gmail.com
 description = Math on (Hyper-Dual) Tensors with Trailing Axes

tensortrax/_tensor.py

@@ -152,6 +152,9 @@ def T(self):
     def ravel(self, order="C"):
         return ravel(self, order=order)
 
+    def reshape(self, *shape, order="C"):
+        return reshape(self, newshape=shape, order=order)
+
     def __matmul__(self, B):
         return matmul(self, B)
 
@@ -222,6 +225,22 @@ def ravel(A, order="C"):
         return np.ravel(A, order=order)
 
 
+def reshape(A, newshape, order="C"):
+    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),
+            ntrax=A.ntrax,
+        )
+    else:
+        return np.reshape(A, newshape=newshape, order=order)
+
+
 def einsum3(subscripts, *operands):
     "Einsum with three operands."
     A, B, C = operands
@@ -387,7 +406,7 @@ def transpose(A):
 
 
 def matmul(A, B):
-    ik = "abcdefghijklm"[13-len(A.shape):]
-    kj = "mnopqrstuvwxy"[: len(B.shape)]
-    ij = (ik + kj).replace("m", "")
+    ik = "ik"[2 - len(A.shape) :]
+    kj = "kj"[: len(B.shape)]
+    ij = (ik + kj).replace("k", "")
     return einsum(f"{ik}...,{kj}...->{ij}...", A, B)

tensortrax/math/__init__.py

@@ -29,5 +29,6 @@
     matmul,
     diagonal,
     ravel,
+    reshape,
 )
 from . import _math_array as array

tensortrax/math/_math_tensor.py

@@ -10,7 +10,7 @@
 
 import numpy as np
 
-from .._tensor import Tensor, ravel, einsum, matmul, f, δ, Δ, Δδ
+from .._tensor import Tensor, ravel, reshape, einsum, matmul, f, δ, Δ, Δδ
 from ._linalg import _linalg_array as array
 
 

tests/test_math.py

@@ -99,6 +99,14 @@ def test_math():
 
     u[0] = t[0]
 
+    t.reshape(9)
+    t.reshape(3, 3)
+
+    tm.reshape(t, (9,))
+    tm.reshape(t, (3, 3))
+
+    tm.reshape(x, (3, 3, 100))
+
 
 if __name__ == "__main__":
     test_math()