Update for coords with batch dims

Author: mfschubert

docs/notebooks/metalens_challenge.ipynb

@@ -151,7 +151,7 @@
     "\n",
     "ex, ey, ez = aux[\"efield\"]\n",
     "x, _, z = aux[\"field_coordinates\"]\n",
-    "xplot, zplot = onp.meshgrid(x[:, 0], z, indexing=\"ij\")\n",
+    "xplot, zplot = onp.meshgrid(x[0, :, 0], z, indexing=\"ij\")\n",
     "\n",
     "abs_field = onp.sqrt(onp.abs(ex) ** 2 + onp.abs(ey) ** 2 + onp.abs(ez) ** 2)\n",
     "\n",