Update line search in LBFGS to use jit and unroll from LBFGS

HajimeKawahara · Feb 19, 2022 · 371462b · 371462b
1 parent 0ff6be8
commit 371462b
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/jaxopt/_src/lbfgs.py b/jaxopt/_src/lbfgs.py
@@ -259,7 +259,9 @@ def update(self,
                                 value_and_grad=True,
                                 maxiter=self.maxls,
                                 decrease_factor=self.decrease_factor,
-                                condition=self.condition)
+                                condition=self.condition,
+                                jit=self.jit,
+                                unroll=self.unroll)
     init_stepsize = state.stepsize * self.increase_factor
     new_stepsize, ls_state = ls.run(init_stepsize=init_stepsize,
                                     params=params, value=value, grad=grad,