adds get_diag, frobenius norm and trace to PMatQuasiDiag

nngeometry/object/pspace.py

@@ -210,6 +210,15 @@ def get_dense_tensor(self):
     def get_diag(self):
         return self.data
 
+    def solve(self, v, regul=1e-8):
+        """
+        solves v = Ax in x
+        """
+        # TODO: test
+        solution = v.to_flat_representation() / (self.data + regul)
+        return PVector(layer_collection=v.layer_collection,
+                       vector_repr=solution)
+
     def __add__(self, other):
         sum_diags = self.data + other.data
         return PMatDiag(generator=self.generator,
@@ -683,10 +692,18 @@ def get_dense_tensor(self):
         return M
 
     def frobenius_norm(self):
-        raise NotImplementedError
+        norm2 = 0
+        for layer_id in self.generator.layer_collection.layers.keys():
+            diag, cross = self.data[layer_id]
+            norm2 += torch.dot(diag, diag)
+            if cross is not None:
+                norm2 += 2 * torch.dot(cross.view(-1), cross.view(-1))
+
+        return norm2 ** .5
 
     def get_diag(self):
-        raise NotImplementedError
+        return torch.cat([self.data[l_id][0] for l_id in
+                          self.generator.layer_collection.layers.keys()])
 
     def inverse(self):
         raise NotImplementedError
@@ -695,7 +712,8 @@ def mv(self):
         raise NotImplementedError
 
     def trace(self):
-        raise NotImplementedError
+        return sum([self.data[l_id][0].sum() for l_id in
+                    self.generator.layer_collection.layers.keys()])
 
     def vTMv(self):
         raise NotImplementedError

tests/test_jacobian.py

@@ -606,3 +606,21 @@ def test_jacobian_pquasidiag_vs_pdense():
                 # compare upper triangular block with lower triangular one
                 check_tensors(matrix_qd[start:start+sw+sb, start+sw:],
                               matrix_qd[start+sw:, start:start+sw+sb].t())
+
+def test_jacobian_pquasidiag():
+    for get_task in [get_fullyconnect_task]:
+        loader, lc, parameters, model, function, n_output = get_task()
+        model.train()
+        generator = Jacobian(layer_collection=lc,
+                             model=model,
+                             loader=loader,
+                             function=function,
+                             n_output=n_output)
+        PMat_qd = PMatQuasiDiag(generator)
+        dense_tensor = PMat_qd.get_dense_tensor()
+
+        check_tensors(torch.diag(dense_tensor), PMat_qd.get_diag())
+
+        check_ratio(torch.norm(dense_tensor), PMat_qd.frobenius_norm())
+
+        check_ratio(torch.trace(dense_tensor), PMat_qd.trace())