Internal change

PiperOrigin-RevId: 544458648

jaxopt/_src/backtracking_linesearch.py

@@ -149,8 +149,8 @@ def update(
       value: Optional[float] = None,
       grad: Optional[Any] = None,
       descent_direction: Optional[Any] = None,
-      fun_args: list[Any] = [],
-      fun_kwargs: dict[str, Any] = {},
+      fun_args: list = [],
+      fun_kwargs: dict = {},
   ) -> base.LineSearchStep:
     """Performs one iteration of backtracking line search.
 

jaxopt/_src/base.py

@@ -475,8 +475,8 @@ def run(self,
           value: Optional[float] = None,
           grad: Optional[Any] = None,
           descent_direction: Optional[Any] = None,
-          fun_args: list[Any] = [],
-          fun_kwargs: dict[str, Any] = {}) -> LineSearchStep:
+          fun_args: list = [],
+          fun_kwargs: dict = {}) -> LineSearchStep:
 
     return super()._run(init_stepsize, params, value, grad, descent_direction,
                         fun_args, fun_kwargs)

jaxopt/_src/broyden.py

@@ -333,7 +333,7 @@ def ls_fun_with_aux(params, *args, **kwargs):
         new_stepsize, ls_state = ls.run(init_stepsize,
                                         params, value, None,
                                         descent_direction,
-                                        *args, **kwargs)
+                                        fun_args=args, fun_kwargs=kwargs)
         new_value, new_aux = ls_state.aux
         new_params = ls_state.params
       else:

jaxopt/_src/hager_zhang_linesearch.py

@@ -337,8 +337,8 @@ def init_state(
       value: Optional[float] = None,
       grad: Optional[Any] = None,
       descent_direction: Optional[Any] = None,
-      fun_args: list[Any] = [],
-      fun_kwargs: dict[str, Any] = {},
+      fun_args: list = [],
+      fun_kwargs: dict = {},
   ) -> HagerZhangLineSearchState:
     """Initialize the line search state.
 
@@ -419,8 +419,8 @@ def update(
       value: Optional[float] = None,
       grad: Optional[Any] = None,
       descent_direction: Optional[Any] = None,
-      fun_args: list[Any] = [],
-      fun_kwargs: dict[str, Any] = {},
+      fun_args: list = [],
+      fun_kwargs: dict = {},
   ) -> base.LineSearchStep:
     """Performs one iteration of Hager-Zhang line search.
 

jaxopt/_src/zoom_linesearch.py

@@ -592,8 +592,8 @@ def init_state(
       value: Optional[float] = None,
       grad: Optional[Any] = None,
       descent_direction: Optional[Any] = None,
-      fun_args: list[Any] = [],
-      fun_kwargs: dict[str, Any] = {},
+      fun_args: list = [],
+      fun_kwargs: dict = {},
   ) -> base.LineSearchStep:
     """Initialize the line search state by computing all relevant quantities and store it in the initial state.
 
@@ -682,8 +682,8 @@ def update(
       value: Optional[float] = None,
       grad: Optional[Any] = None,
       descent_direction: Optional[Any] = None,
-      fun_args: list[Any] = [],
-      fun_kwargs: dict[str, Any] = {},
+      fun_args: list = [],
+      fun_kwargs: dict = {},
   ) -> base.LineSearchStep:
     """Combines Algorithms 3.5 and 3.6 of [1].