Fix introductory CNN example in recent JAX.

examples/mnist.ipynb

@@ -594,7 +594,9 @@
     "        y: Int[Array, \" batch\"],\n",
     "    ):\n",
     "        loss_value, grads = eqx.filter_value_and_grad(loss)(model, x, y)\n",
-    "        updates, opt_state = optim.update(grads, opt_state, model)\n",
+    "        updates, opt_state = optim.update(\n",
+    "            grads, opt_state, eqx.filter(model, eqx.is_array)\n",
+    "        )\n",
     "        model = eqx.apply_updates(model, updates)\n",
     "        return model, opt_state, loss_value\n",
     "\n",