Fix nnx guide error (#2183)

* Update keras_nnx_guide.py

* Update keras_nnx_guide.ipynb

* Update keras_nnx_guide.md

Author: divyashreepathihalli

guides/ipynb/keras_nnx_guide.ipynb

@@ -342,7 +342,8 @@
     "        y_pred = model_(x)\n",
     "        return jnp.mean((y - y_pred) ** 2)\n",
     "\n",
-    "    grads = nnx.grad(loss_fn, wrt=trainable_var)(model)\n",
+    "    diff_state = nnx.DiffState(0, trainable_var)\n",
+    "    grads = nnx.grad(loss_fn, argnums=diff_state)(model)\n",
     "    optimizer.update(model, grads)\n",
     "\n",
     "\n",
@@ -506,4 +507,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 0
-}
+}

guides/keras_nnx_guide.py

@@ -208,7 +208,8 @@ def loss_fn(model_):
         y_pred = model_(x)
         return jnp.mean((y - y_pred) ** 2)
 
-    grads = nnx.grad(loss_fn, wrt=trainable_var)(model)
+    diff_state = nnx.DiffState(0, trainable_var)
+    grads = nnx.grad(loss_fn, argnums=diff_state)(model)
     optimizer.update(model, grads)
 
 

guides/md/keras_nnx_guide.md

@@ -223,7 +223,8 @@ def train_step(model, optimizer, batch):
         y_pred = model_(x)
         return jnp.mean((y - y_pred) ** 2)
 
-    grads = nnx.grad(loss_fn, wrt=trainable_var)(model)
+    diff_state = nnx.DiffState(0, trainable_var)
+    grads = nnx.grad(loss_fn, argnums=diff_state)(model)
     optimizer.update(model, grads)