Merge pull request #86 from tfjgeorge/pvec_pow

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()