Add flax.nnx.value_and_grad docstring

[8bitmp3]

@cgarciae @IvyZX @levskaya

[cgarciae]

@8bitmp3 some tests are failing.

[8bitmp3]

@8bitmp3 some tests are failing.

Found it. ... was missing:

>>> value_and_grad_fn = nnx.value_and_grad(loss_fn)
...
>>> values, grads = value_and_grad_fn(m, x, y)

[cgarciae]

@8bitmp3 CI is still failing

[8bitmp3]

The error:

367     >>> values, grads = value_and_grad_fn(m, x, y)
368     >>> print(jax.tree.map(jnp.shape, grads))
Expected nothing
Got:
    State({
      'bias': VariableState(
        type=Param,
        value=(3,)
      ),
      'kernel': VariableState(
        type=Param,
        value=(2, 3)
      )
    })

I could try without print() @cgarciae

[cgarciae]

@8bitmp3 copy the output you see there and paste it on the example e.g:

>>> print(jax.tree.map(jnp.shape, grads))
State({
  'bias': VariableState(
    type=Param,
    value=(3,)
  ),
  'kernel': VariableState(
    type=Param,
    value=(2, 3)
  )
})