No public description

PiperOrigin-RevId: 672160611
google · Sep 9, 2024 · 78b295b · 78b295b
1 parent 6d88292
commit 78b295b
Showing 1 changed file with 3 additions and 5 deletions.
diff --git a/jaxopt/_src/gradient_descent.py b/jaxopt/_src/gradient_descent.py
@@ -79,11 +79,9 @@ def init_state(self,
     """
     return super().init_state(init_params, None, *args, **kwargs)
 
-  def update(self,
-             params: Any,
-             state: NamedTuple,
-             *args,
-             **kwargs) -> base.OptStep:
+  def update(
+      self, params: Any, state: ProxGradState, *args, **kwargs
+  ) -> base.OptStep:
     """Performs one iteration of gradient descent.
 
     Args: