Repository: google/flax

Issue:
Support for optax lbfgs and related optimizers with NNX
I am trying to use L-BFGS and related optimizers with nnx + optax, but running into trouble.  It might be that `optax` has a slightly different optimization interface in those cases: https://optax.readthedocs.io/en/latest/api/optimizers.html#lbfgs with a function called `optax.value_and_grad_from_state` and also a change to the `optimizer.update` interface?

In particular, note that the sample code for these optimizers in https://optax.readthedocs.io/en/latest/api/optimizers.html#lbfgs looks like
```python
opt_state = solver.init(params)
value_and_grad = optax.value_and_grad_from_state(f)
for _ in range(5):
  value, grad = value_and_grad(params, state=opt_state)
  updates, opt_state = solver.update(
     grad, opt_state, params, value=value, grad=grad, value_fn=f
  )
  params = optax.apply_updates(params, updates)
```

So maybe we need the ability to pass a `value_fn` argument on?  Any easy fixes or things I am missing?

### Problem you have encountered:
Take the following implementation LLS with NNX + optax
```python
import jax
import jax.numpy as jnp
from jax import random
import optax
import equinox as eqx
import jax_dataloader as jdl
from jax_dataloader.loaders import DataLoaderJAX
from flax import nnx

N = 64  # samples
M = 2
sigma = 0.001
rngs = nnx.Rngs(42)
theta = random.normal(rngs(), (M,))
X = random.normal(rngs(), (N, M))
Y = X @ theta + sigma * random.normal(rngs(), (N,))  # Adding noise

def residual(model, x, y):
    y_hat = model(x)
    return (y_hat - y) ** 2

def residuals_loss(model, X, Y):
    return jnp.mean(jax.vmap(residual, in_axes=(None, 0, 0))(model, X, Y))

model = nnx.Linear(M, 1, use_bias=False, rngs=rngs)

lr = 0.001
optimizer = nnx.Optimizer(model,
                          optax.sgd(lr) #optax.lbfgs()
                          )

@nnx.jit
def train_step(model, optimizer, X, Y):
    grad_fn = nnx.value_and_grad(residuals_loss, has_aux=False)
    loss, grads = grad_fn(model, X, Y)
    optimizer.update(grads)
    return loss

num_epochs = 500
batch_size = 64
dataset = jdl.ArrayDataset(X, Y)
train_loader = DataLoaderJAX(dataset, batch_size=batch_size, shuffle=True)
for epoch in range(num_epochs):
    for X_batch, Y_batch in train_loader:
        loss = train_step(model, optimizer, X_batch, Y_batch)

    if epoch % 100 == 0:
        print(
            f"Epoch {epoch},||theta - theta_hat|| = {jnp.linalg.norm(theta - jnp.squeeze(model.kernel.value))}"
        )

print(f"||theta - theta_hat|| = {jnp.linalg.norm(theta - jnp.squeeze(model.kernel.value))}")
```

But if I change the optimizer from `optax.sgd(lr)` to be `optax.lbfgs()` then I would expect NNX to work.

### Logs, error messages, etc:
The error it gives is
```python
  File "/Users/jlperla/Documents/GitHub/ECON622_instructor/lectures/examples/linear_regression_jax_nnx.py", line 52, in <module>
    loss = train_step(model, optimizer, X_batch, Y_batch)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/jlperla/anaconda3/envs/econ622/lib/python3.11/site-packages/flax/nnx/nnx/graph.py", line 1043, in update_context_manager_wrapper
    return f(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^
  File "/Users/jlperla/anaconda3/envs/econ622/lib/python3.11/site-packages/flax/nnx/nnx/transforms/transforms.py", line 359, in jit_wrapper
    out, output_state, output_graphdef = jitted_fn(
                                         ^^^^^^^^^^
  File "/Users/jlperla/anaconda3/envs/econ622/lib/python3.11/site-packages/flax/nnx/nnx/transforms/transforms.py", line 158, in jit_fn
    out = f(*args, **kwargs)
          ^^^^^^^^^^^^^^^^^^
  File "/Users/jlperla/Documents/GitHub/ECON622_instructor/lectures/examples/linear_regression_jax_nnx.py", line 43, in train_step
    optimizer.update(grads)
  File "/Users/jlperla/anaconda3/envs/econ622/lib/python3.11/site-packages/flax/nnx/nnx/training/optimizer.py", line 201, in update
    updates, new_opt_state = self.tx.update(grads, self.opt_state, state)
                             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/jlperla/anaconda3/envs/econ622/lib/python3.11/site-packages/optax/transforms/_combining.py", line 73, in update_fn
    updates, new_s = fn(updates, s, params, **extra_args)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: scale_by_zoom_linesearch.<locals>.update_fn() missing 3 required keyword-only arguments: 'value', 'grad', and 'value_fn'
```

### Steps to reproduce:
Run the example above with the optimizer swapped out, i.e.
```python
import jax
import jax.numpy as jnp
from jax import random
import optax
import equinox as eqx
import jax_dataloader as jdl
from jax_dataloader.loaders import DataLoaderJAX
from flax import nnx

N = 64  # samples
M = 2
sigma = 0.001
rngs = nnx.Rngs(42)
theta = random.normal(rngs(), (M,))
X = random.normal(rngs(), (N, M))
Y = X @ theta + sigma * random.normal(rngs(), (N,))  # Adding noise

def residual(model, x, y):
    y_hat = model(x)
    return (y_hat - y) ** 2

def residuals_loss(model, X, Y):
    return jnp.mean(jax.vmap(residual, in_axes=(None, 0, 0))(model, X, Y))

model = nnx.Linear(M, 1, use_bias=False, rngs=rngs)

lr = 0.001
optimizer = nnx.Optimizer(model,
                          optax.lbfgs()
                          )

@nnx.jit
def train_step(model, optimizer, X, Y):
    grad_fn = nnx.value_and_grad(residuals_loss, has_aux=False)
    loss, grads = grad_fn(model, X, Y)
    optimizer.update(grads)
    return loss

num_epochs = 500
batch_size = 64
dataset = jdl.ArrayDataset(X, Y)
train_loader = DataLoaderJAX(dataset, batch_size=batch_size, shuffle=True)
for epoch in range(num_epochs):
    for X_batch, Y_batch in train_loader:
        loss = train_step(model, optimizer, X_batch, Y_batch)

    if epoch % 100 == 0:
        print(
            f"Epoch {epoch},||theta - theta_hat|| = {jnp.linalg.norm(theta - jnp.squeeze(model.kernel.value))}"
        )

print(f"||theta - theta_hat|| = {jnp.linalg.norm(theta - jnp.squeeze(model.kernel.value))}")
```

Discussion:

Patch:
diff --git a/flax/nnx/training/optimizer.py b/flax/nnx/training/optimizer.py
index fc3b4eeb..4b85d5a3 100644
--- a/flax/nnx/training/optimizer.py
+++ b/flax/nnx/training/optimizer.py
@@ -193,7 +193,7 @@ class Optimizer(Object):
     self.opt_state = _wrap_optimizer_state(tx.init(nnx.state(model, wrt)))
     self.wrt = wrt
 
-  def update(self, grads):
+  def update(self, grads, **kwargs):
     """Updates ``step``, ``params``, ``opt_state`` and ``**kwargs`` in return value.
     The ``grads`` must be derived from ``nnx.grad(..., wrt=self.wrt)``, where the
     gradients are with respect to the same :class:`Variable` types as defined in
@@ -249,14 +249,16 @@ class Optimizer(Object):
 
     Args:
       grads: the gradients derived from ``nnx.grad``.
+      **kwargs: additional keyword arguments passed to the tx.update, to support
+      ``GradientTransformationExtraArgs``, such as ``optax.scale_by_backtracking_linesearch``.
     """
     params = nnx.state(self.model, self.wrt)
     opt_state = _opt_state_variables_to_state(self.opt_state)
 
-    updates, new_opt_state = self.tx.update(grads, opt_state, params)
+    updates, new_opt_state = self.tx.update(grads, opt_state, params, **kwargs)
     new_params = optax.apply_updates(params, updates)
     assert isinstance(new_params, nnx.State)
 
     self.step.value += 1
     nnx.update(self.model, new_params)
-    _update_opt_state(self.opt_state, new_opt_state)
+    _update_opt_state(self.opt_state, new_opt_state)
\ No newline at end of file


Test patch:
diff --git a/tests/nnx/optimizer_test.py b/tests/nnx/optimizer_test.py
index c74160cc..5c28e727 100644
--- a/tests/nnx/optimizer_test.py
+++ b/tests/nnx/optimizer_test.py
@@ -128,6 +128,58 @@ class TestOptimizer(parameterized.TestCase):
 
     self.assertTrue(new_loss < initial_loss)
 
+
+  @parameterized.product(
+    module_cls=[nnx.Linear, Model],
+    jit_decorator=[lambda f: f, nnx.jit, jax.jit],
+    optimizer=[optax.lbfgs],
+  )
+  def test_jit_linesearch(self, module_cls, jit_decorator, optimizer):
+    x = jax.random.normal(jax.random.key(0), (1, 2))
+    y = jnp.ones((1, 4))
+    model = module_cls(2, 4, rngs=nnx.Rngs(0))
+    tx = optimizer(
+      1e-3
+    )
+    state = nnx.Optimizer(model, tx)
+
+    if jit_decorator == jax.jit:
+      model_static, model_state = nnx.split(state.model)
+      loss_fn = lambda graphdef, state, x, y: (
+        (nnx.merge(graphdef, state)(x) - y) ** 2
+      ).mean()
+      initial_loss = loss_fn(model_static, model_state, x, y)
+
+      def jax_jit_train_step(graphdef, state, x, y):
+        state = nnx.merge(graphdef, state)
+        model_static, model_state = nnx.split(state.model)
+        grads = jax.grad(loss_fn, argnums=1)(model_static, model_state, x, y)
+        state.update(grads, grad = grads, value = initial_loss, value_fn = lambda state: loss_fn(model_static, state, x, y))
+        return nnx.split(state)
+
+      graphdef, state = jit_decorator(jax_jit_train_step)(
+        *nnx.split(state), x, y
+      )
+      state = nnx.merge(graphdef, state)
+      new_loss = loss_fn(*nnx.split(state.model), x, y)
+
+    else:
+      graphdef = nnx.graphdef(model)
+      loss_fn = lambda model, x, y: ((model(x) - y) ** 2).mean()
+
+      loss_fn_split = lambda state: loss_fn(nnx.merge(graphdef, state), x, y)
+
+      initial_loss = loss_fn(state.model, x, y)
+
+      def nnx_jit_train_step(optimizer: nnx.Optimizer, x, y):
+        grads = nnx.grad(loss_fn)(optimizer.model, x, y)
+        optimizer.update(grads, grad = grads, value = initial_loss, value_fn = loss_fn_split)
+
+      jit_decorator(nnx_jit_train_step)(state, x, y)
+      new_loss = loss_fn(state.model, x, y)
+
+    self.assertTrue(new_loss < initial_loss)
+
   @parameterized.product(
     module_cls=[nnx.Linear, Model],
     optimizer=[optax.sgd, optax.adam],
@@ -203,6 +255,55 @@ class TestOptimizer(parameterized.TestCase):
         )
       )
 
+  @parameterized.parameters(
+    {'variable': nnx.Param},
+    #{'variable': nnx.LoRAParam},
+    {'variable': (nnx.Param, nnx.LoRAParam)},
+  )
+  def test_wrt_update_linesearch(self, variable):
+    in_features = 4
+    out_features = 10
+    model = nnx.LoRA(
+      in_features=in_features,
+      lora_rank=2,
+      out_features=out_features,
+      base_module=Model(
+        in_features=in_features, out_features=out_features, rngs=nnx.Rngs(0)
+      ),
+      rngs=nnx.Rngs(1),
+    )
+    state = nnx.Optimizer(model, optax.lbfgs(), wrt=variable)
+    prev_variables, prev_other_variables = nnx.state(model, variable, ...)
+
+    x = jnp.ones((1, 4))
+    y = jnp.ones((1, 10))
+    loss_fn = lambda model, x, y: ((model(x) - y) ** 2).mean()
+
+    grads = nnx.grad(loss_fn, argnums=nnx.DiffState(0, variable))(
+      state.model, x, y
+    )
+    initial_loss = loss_fn(model, x, y)
+    graphdef = nnx.graphdef(model)
+    loss_fn_split = lambda state: loss_fn(nnx.merge(graphdef, state), x, y)
+
+    state.update(grads, grad=grads, value_fn = loss_fn_split, value = initial_loss)
+    self.assertTrue(loss_fn(model, x, y) < initial_loss)
+
+    # make sure only the Variable's filtered in `wrt` are changed, and the others are unchanged
+    variables, other_variables = nnx.state(model, variable, ...)
+    self.assertTrue(
+      jax.tree.all(
+        jax.tree.map(lambda x, y: (x != y).all(), prev_variables, variables)
+      )
+    )
+    if other_variables:
+      self.assertTrue(
+        jax.tree.all(
+          jax.tree.map(
+            lambda x, y: (x == y).all(), prev_other_variables, other_variables
+          )
+        )
+      )
 
 if __name__ == '__main__':
   absltest.main()