Fix typo

Author: samanklesaria

docs_nnx/guides/Optimization Cookbook.ipynb

@@ -59,9 +59,6 @@
     "        nnx.Linear(2,8, rngs=rngs, kernel_init=ft.partial(param_init, **kwargs)),\n",
     "        nnx.Linear(8,8, rngs=rngs, kernel_init=ft.partial(param_init, **kwargs)))\n",
     "\n",
-    "def nnx_loss_fn(params, x, y):\n",
-    "    return jnp.sum((y - model(params, x))**2)\n",
-    "\n",
     "def nnx_loss_fn(model, x, y):\n",
     "    return jnp.sum((model(x) - y) ** 2)"
    ]

docs_nnx/guides/Optimization Cookbook.md

@@ -6,6 +6,10 @@ jupyter:
       format_name: markdown
       format_version: '1.3'
       jupytext_version: 1.13.8
+  kernelspec:
+    display_name: Python 3 (ipykernel)
+    language: python
+    name: python3
 ---
 
 # A Flax Optimization Cookbook
@@ -38,9 +42,6 @@ def nnx_model(rngs, **kwargs):
         nnx.Linear(2,8, rngs=rngs, kernel_init=ft.partial(param_init, **kwargs)),
         nnx.Linear(8,8, rngs=rngs, kernel_init=ft.partial(param_init, **kwargs)))
 
-def nnx_loss_fn(params, x, y):
-    return jnp.sum((y - model(params, x))**2)
-
 def nnx_loss_fn(model, x, y):
     return jnp.sum((model(x) - y) ** 2)
 ```