Merge pull request #18 from adtzlr/fix-set-items

Fix `Tensor.__setitems__()`

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "tensortrax"
-version = "0.1.7"
+version = "0.1.8"
 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.1.7
+version = 0.1.8
 author = Andreas Dutzler
 author_email = a.dutzler@gmail.com
 description = Math on (Hyper-Dual) Tensors with Trailing Axes

tensortrax/_tensor.py

@@ -168,6 +168,18 @@ def __getitem__(self, key):
     def __setitem__(self, key, value):
         if isinstance(value, Tensor):
             self.x[key] = f(value)
+            if self.δx[key].shape != δ(value).shape:
+                self.δx = transpose(
+                    np.resize(self.δx, (*self.trax[::-1], *self.shape[::-1])).T
+                )
+            if self.Δx[key].shape != Δ(value).shape:
+                self.Δx = transpose(
+                    np.resize(self.Δx, (*self.trax[::-1], *self.shape[::-1])).T
+                )
+            if self.Δδx[key].shape != Δδ(value).shape:
+                self.Δδx = transpose(
+                    np.resize(self.Δδx, (*self.trax[::-1], *self.shape[::-1])).T
+                )
             self.δx[key] = δ(value)
             self.Δx[key] = Δ(value)
             self.Δδx[key] = Δδ(value)

tests/test_math.py

@@ -93,6 +93,12 @@ def test_math():
     T[:, 0] = F[:, 0]
     T[:, 0] = T[:, 0]
 
+    x = np.ones((3, 3, 100))
+    t = tr.Tensor(x, x, x, x, ntrax=1)
+    u = tr.Tensor(x, ntrax=1)
+
+    u[0] = t[0]
+
 
 if __name__ == "__main__":
     test_math()