No public description

PiperOrigin-RevId: 672706385
google · Sep 17, 2024 · a9b974f · a9b974f
1 parent 66d90db
commit a9b974f
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: