solveJ dense

Author: tfjgeorge

nngeometry/object/pspace.py

@@ -164,11 +164,14 @@ def solveJ(self, J, regul=1e-8):
         """
         solves J = AX in X
         """
+        J_torch = J.to_torch()
+        sJ = J_torch.size()
         inv_v = torch.linalg.solve(
             self.data + regul * torch.eye(self.size(0), device=self.data.device),
-            J.to_torch()[0].t(),
+            J_torch.view(-1, sJ[-1]),
+            left=False,
         )
-        return inv_v
+        return PFMapDense(generator=self.generator, data=inv_v.reshape(*sJ))
 
     def inverse(self, regul=1e-8):
         inv_tensor = torch.inverse(
@@ -667,7 +670,9 @@ def to_torch(self, split_weight_bias=True):
         """
         _, diags = self.data
         s = self.generator.layer_collection.numel()
-        M = torch.zeros((s, s), device=self.generator.get_device(), dtype=self.generator.get_dtype())
+        M = torch.zeros(
+            (s, s), device=self.generator.get_device(), dtype=self.generator.get_dtype()
+        )
         KFE_layers = self.get_KFE(split_weight_bias=split_weight_bias)
         for layer_id, _ in self.generator.layer_collection.layers.items():
             diag = diags[layer_id]

tests/test_torch_hooks.py

@@ -341,19 +341,33 @@ def test_jacobian_pdense():
             # Test solve
             # NB: regul is high since the matrix is not full rank
             regul = 1e-3
-            Mv_regul = torch.mv(
-                PMat_dense.to_torch()
-                + regul * torch.eye(PMat_dense.size(0), device=device),
-                dw.to_torch(),
+            Mv_torch = torch.mv(PMat_dense.to_torch(), dw.to_torch())
+            Mv_regul = PVector(
+                layer_collection=lc, vector_repr=Mv_torch + regul * dw.to_torch()
             )
-            Mv_regul = PVector(layer_collection=lc, vector_repr=Mv_regul)
-            dw_using_inv = PMat_dense.solve(Mv_regul, regul=regul)
+            dw_solve = PMat_dense.solve(Mv_regul, regul=regul)
             check_tensors(
                 dw.to_torch(),
-                dw_using_inv.to_torch(),
+                dw_solve.to_torch(),
                 eps=5e-3,
             )
 
+            # Test solve with jacobian
+            # TODO improve
+            c = 1.678
+            stacked_mv = torch.stack((Mv_torch, c * Mv_torch)).unsqueeze(0)
+            stacked_v = torch.stack((dw.to_torch(), c * dw.to_torch())).unsqueeze(0)
+            jaco = PFMapDense(
+                generator=generator,
+                data=stacked_mv + regul * stacked_v,
+            )
+            J_back = PMat_dense.solveJ(jaco, regul=regul)
+
+            check_tensors(
+                stacked_v,
+                J_back.to_torch(),
+            )
+
             # Test inv
             PMat_inv = PMat_dense.inverse(regul=regul)
             check_tensors(

tests/test_torch_hooks_ekfac.py

@@ -108,6 +108,7 @@ def test_pspace_ekfac_vs_direct():
             check_tensors(v.to_torch(), v_back.to_torch())
 
             # Test solve with jacobian
+            # TODO improve
             c = 1.678
             stacked_mv = torch.stack(
                 (mv_ekfac.to_torch(), c * mv_ekfac.to_torch())