Fix flax test for flax v0.9.0 (#11)

jax_ml_stack/tests/test_nnx_with_optax.py

@@ -46,7 +46,7 @@ def loss(model, x=x, y=y):
       return jnp.mean((model(x) - y) ** 2)
 
     initial_loss = loss(model)
-    grads = nnx.grad(loss, wrt=nnx.Param)(state.model)
+    grads = nnx.grad(loss)(state.model)
     state.update(grads)
     final_loss = loss(model)