minor change to make backtracking init and update return the same sta…

…te types, this avoids recompilation when calling update
google-deepmind · Jan 15, 2025 · 7d090b7 · 7d090b7
1 parent 98c73c5
commit 7d090b7
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/optax/_src/linesearch.py b/optax/_src/linesearch.py
@@ -255,11 +255,11 @@ def init_fn(params: base.Params) -> ScaleByBacktrackingLinesearchState:
       grad = None
     return ScaleByBacktrackingLinesearchState(
         learning_rate=jnp.array(1.0),
-        value=jnp.array(jnp.inf),
+        value=jnp.array(jnp.inf, dtype=params.dtype),
         grad=grad,
         info=BacktrackingLinesearchInfo(
             num_linesearch_steps=0,
-            decrease_error=jnp.array(jnp.inf),
+            decrease_error=jnp.array(jnp.inf, dtype=params.dtype),
         ),
     )