has_aux fix

timothyn617 · KfacJaxDev · commit f9b8b1ba3bda · 2025-05-23T01:37:04.000-07:00
PiperOrigin-RevId: 762317513
diff --git a/examples/optimizers.py b/examples/optimizers.py
@@ -124,7 +124,9 @@ def create_optimizer(
         "Schedule Free is only supported for optax optimizers."
     )
 
-  value_and_grad_func = jax.value_and_grad(train_model_func, has_aux=has_aux)
+  value_and_grad_func = jax.value_and_grad(
+      train_model_func, has_aux=has_aux or has_func_state
+  )
 
   kwargs = dict(**config[name])
 

Original file line number	Diff line number	Diff line change
`@@ -124,7 +124,9 @@ def create_optimizer(`
`124`	`124`	`"Schedule Free is only supported for optax optimizers."`
`125`	`125`	`)`
`126`	`126`
`127`		`- value_and_grad_func = jax.value_and_grad(train_model_func, has_aux=has_aux)`
	`127`	`+ value_and_grad_func = jax.value_and_grad(`
	`128`	`+ train_model_func, has_aux=has_aux or has_func_state`
	`129`	`+ )`
`128`	`130`
`129`	`131`	`kwargs = dict(**config[name])`
`130`	`132`