Merge pull request #785 from yixiaoer:type_annotation

PiperOrigin-RevId: 605532129

docs/optax-101.ipynb

@@ -306,7 +306,7 @@
         "}\n",
         "\n",
         "\n",
-        "def net(x: jnp.ndarray, params: jnp.ndarray) -\u003e jnp.ndarray:\n",
+        "def net(x: jnp.ndarray, params: optax.Params) -\u003e jnp.ndarray:\n",
         "  x = jnp.dot(x, params['hidden'])\n",
         "  x = jax.nn.relu(x)\n",
         "  x = jnp.dot(x, params['output'])\n",