Remove batch_flat

backpack/core/derivatives/basederivatives.py

@@ -129,12 +129,6 @@ def make_residual_mat_prod(self, module, g_inp, g_out):
         """
         raise NotImplementedError
 
-    # TODO Refactor and remove
-    def batch_flat(self, tensor):
-        batch = tensor.size(0)
-        # TODO Removing the clone().detach() will destroy the computation graph
-        # Tests will fail
-        return batch, tensor.clone().detach().view(batch, -1)
 
     # TODO Refactor and remove
     def get_batch(self, module):

backpack/core/derivatives/elementwise.py

@@ -18,8 +18,9 @@ def _jac_mat_prod(self, module, g_inp, g_out, mat):
     def ea_jac_t_mat_jac_prod(self, module, g_inp, g_out, mat):
         self._no_inplace(module)
 
-        batch, df_flat = self.batch_flat(self.df(module, g_inp, g_out))
-        return einsum("ni,nj,ij->ij", (df_flat, df_flat, mat)) / batch
+        N = module.input0.size(0)
+        df_flat = self.df(module, g_inp, g_out).reshape(N, -1)
+        return einsum("ni,nj,ij->ij", (df_flat, df_flat, mat)) / N
 
     def hessian_diagonal(self, module, g_inp, g_out):
         self._no_inplace(module)

backpack/extensions/secondorder/hbp/linear.py

@@ -55,9 +55,9 @@ def _bias_for_batch_average(self, backproped):
         return [backproped]
 
     def __mean_input(self, module):
-        _, flat_input = self.derivatives.batch_flat(module.input0)
-        return flat_input.mean(0)
+        return module.input0.mean(0).flatten()
 
     def __mean_input_outer(self, module):
-        N, flat_input = self.derivatives.batch_flat(module.input0)
+        N = module.input0.size(0)
+        flat_input = module.input0.reshape(N, -1)
         return einsum("ni,nj->ij", (flat_input, flat_input)) / N