No public description

PiperOrigin-RevId: 672160611

Author: JAXopt authors

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: