Merge pull request #92 from BenjaminDev/91_Laplace_GGN_arguments

Laplace ggn arguments

posteriors/laplace/dense_ggn.py

@@ -137,8 +137,8 @@ def outer_loss(z, batch):
 
     with torch.no_grad(), CatchAuxError():
         ggn_batch, aux = ggn(
-            partial(forward, batch=batch),
-            partial(outer_loss, batch=batch),
+            lambda params: forward(params, batch),
+            lambda z: outer_loss(z, batch),
             forward_has_aux=True,
             loss_has_aux=False,
             normalize=False,

posteriors/laplace/diag_ggn.py

@@ -134,8 +134,8 @@ def outer_loss(z, batch):
 
     with torch.no_grad(), CatchAuxError():
         diag_ggn_batch, aux = diag_ggn(
-            partial(forward, batch=batch),
-            partial(outer_loss, batch=batch),
+            lambda params: forward(params, batch),
+            lambda z: outer_loss(z, batch),
             forward_has_aux=True,
             loss_has_aux=False,
             normalize=False,

tests/laplace/test_dense_ggn.py

@@ -18,8 +18,8 @@ def normal_log_likelihood(y_pred, batch):
     )  # validate args introduces control flows not yet supported in torch.func.vmap
 
 
-def forward_m(params, batch, model):
-    y_pred = functional_call(model, params, batch[0])
+def forward_m(params, b, model):
+    y_pred = functional_call(model, params, b[0])
     return y_pred, torch.tensor([])
 
 

tests/laplace/test_diag_ggn.py

@@ -18,8 +18,8 @@ def normal_log_likelihood(y_pred, batch):
     )  # validate args introduces control flows not yet supported in torch.func.vmap
 
 
-def forward_m(params, batch, model):
-    y_pred = functional_call(model, params, batch[0])
+def forward_m(params, b, model):
+    y_pred = functional_call(model, params, b[0])
     return y_pred, torch.tensor([])