Fix kwargs parsing

Author: mcbal

deep_implicit_attention/deq.py

@@ -65,7 +65,7 @@ def _fixed_point(self, z0, x, *args, **kwargs):
             out = self.solver(
                 lambda z: self.fun(z, x, *args),
                 z0,
-                **filter_kwargs(kwargs, "solver_fwd"),
+                **filter_kwargs(kwargs, "solver_fwd_"),
             )
             z, _ = out["result"], out["rel_trace"]
 
@@ -81,7 +81,7 @@ def backward_hook(grad):
                     lambda y: autograd.grad(fun_bwd, z_bwd, y, retain_graph=True)[0]
                     + grad,
                     torch.zeros_like(grad),
-                    **filter_kwargs(kwargs, "solver_bwd"),
+                    **filter_kwargs(kwargs, "solver_bwd_"),
                 )
                 g, _ = out["result"], out["rel_trace"]
                 return g
@@ -94,7 +94,7 @@ def forward(self, x, *args, **kwargs):
         # Get list of initial guess tensors and reshape into a batch of vectors
         z0 = self.fun.pack_state(kwargs.get("z0", self.fun.get_initial_guess(x)))
         # Find equilibrium vectors
-        self.kwargs.update(**kwargs)
-        z_star = self._fixed_point(z0, x, *args, **kwargs)
+        z_star = self._fixed_point(z0, x, *args, **self.kwargs)
         # Return (subset of) list of tensors of original input shapes
-        return [self.fun.unpack_state(z_star)[i] for i in self.output_elements]
+        out = [self.fun.unpack_state(z_star)[i] for i in self.output_elements]
+        return out[0] if len(out) == 1 else out

deep_implicit_attention/modules.py

@@ -71,8 +71,8 @@ def __init__(
         self.prior_init_std = prior_init_std
 
         self.solver = solver
-        self.solver_tol = 1e-4
-        self.solver_max_iter = 30
+        self.solver_tol = solver_tol
+        self.solver_max_iter = solver_max_iter
 
         self.lin_response_correction = lin_response_correction
 
@@ -144,4 +144,3 @@ def forward(self, z, x, *args):
         next_spin_mean, next_spin_var = pf * (cav_mean + x), pf
 
         return self.pack_state([next_spin_mean, next_spin_var])
-

deep_implicit_attention/utils.py

@@ -12,8 +12,8 @@ def set_all_seeds(seed):
     torch.cuda.manual_seed_all(seed)
 
 
-def filter_kwargs(kwargs, str):
-    return {k.replace(str, ""): v for k, v in kwargs if k.startswith(str)}
+def filter_kwargs(kwargs, prefix):
+    return {k.replace(prefix, ""): v for k, v in kwargs.items() if k.startswith(prefix)}
 
 
 def make_traceless(X: torch.Tensor):