pvector pow

nngeometry/object/vector.py

@@ -295,6 +295,21 @@ def __sub__(self, other):
                 ),
             )
 
+    def __pow__(self, exp):
+        if self.dict_repr is not None:
+            v_dict = dict()
+            for l_id, l in self.layer_collection.layers.items():
+                if l.bias is not None:
+                    v_dict[l_id] = (
+                        self.dict_repr[l_id][0] ** exp,
+                        self.dict_repr[l_id][1] ** exp,
+                    )
+                else:
+                    v_dict[l_id] = (self.dict_repr[l_id][0] ** exp,)
+            return PVector(self.layer_collection, dict_repr=v_dict)
+        else:
+            return PVector(self.layer_collection, vector_repr=self.vector_repr**exp)
+
     def dot(self, other):
         """
         Computes the dot product between `self` and `other`

requirements.txt

@@ -1,3 +1,3 @@
-torch==2.0.1
+torch>=2.0.1
 torchvision>=0.9.1
 requests>=2.24.0

tests/test_vector.py

@@ -122,6 +122,28 @@ def test_sub():
     )
 
 
+def test_pow():
+    model = ConvNet()
+    layer_collection = LayerCollection.from_model(model)
+    r1 = random_pvector(layer_collection)
+    sqrt_r1 = r1**3
+    assert (
+        torch.norm(
+            sqrt_r1.get_flat_representation() - r1.get_flat_representation() ** 3
+        )
+        < 1e-5
+    )
+
+    r1 = random_pvector_dict(layer_collection)
+    sqrt_r1 = r1**3
+    assert (
+        torch.norm(
+            sqrt_r1.get_flat_representation() - r1.get_flat_representation() ** 3
+        )
+        < 1e-5
+    )
+
+
 def test_clone():
     eps = 1e-8
     model = ConvNet()