Reuse the vectorize=True example for vectorize=[dims]

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Author: headtr1ck

advanced/apply_ufunc/automatic-vectorizing-numpy.ipynb

@@ -344,9 +344,7 @@
    "source": [
     "## Looping over some dimensions only\n",
     "\n",
-    "`vectorize=True` loops over all loop dimensions, so `debug_interp` was called 12 times. But many functions can process a whole stack of vectors at once, as long as the other inputs are the same for all of them. For example, `scipy.interpolate.make_interp_spline` interpolates all vectors along the last axis at once, but only to a single set of new latitudes.\n",
-    "\n",
-    "Let's use different new latitudes at every time step. They vary along `time`, but not along `lon`, so we only need to loop over `time`. Instead of `True`, we pass the dimensions to loop over to `vectorize`. All other loop dimensions, here `lon`, are passed to our function in bulk:"
+    "Let's use different new latitudes at every time step. They now vary along `time`, but not along `lon`. The same call as above still works:"
    ]
   },
   {
@@ -358,12 +356,45 @@
    },
    "outputs": [],
    "source": [
-    "from scipy.interpolate import make_interp_spline\n",
-    "\n",
     "newlat_t = xr.DataArray(\n",
     "    np.linspace(20, 70, 100) + np.arange(4)[:, np.newaxis], dims=(\"time\", \"lat\")\n",
     ")\n",
     "\n",
+    "interped = xr.apply_ufunc(\n",
+    "    debug_interp,  # first the function\n",
+    "    newlat_t,\n",
+    "    air.lat,\n",
+    "    air,\n",
+    "    input_core_dims=[[\"lat\"], [\"lat\"], [\"lat\"]],\n",
+    "    output_core_dims=[[\"lat\"]],\n",
+    "    exclude_dims={\"lat\"},  # dimensions allowed to change size. Must be set!\n",
+    "    vectorize=True,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "16",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "`debug_interp` is again called 12 times, but the 3 calls for the `lon` values of each `time` all use the same new latitudes. Many functions can process a whole stack of vectors at once in that case. For example, `scipy.interpolate.make_interp_spline` interpolates all vectors along the last axis at once, to a single set of new latitudes.\n",
+    "\n",
+    "So we only need to loop over `time`. Instead of `True`, we pass the dimensions to loop over to `vectorize`. All other loop dimensions, here `lon`, are passed to our function in bulk:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "17",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "from scipy.interpolate import make_interp_spline\n",
+    "\n",
     "\n",
     "def debug_interp_batched(xi, x, data):\n",
     "    print(f\"data: {data.shape} | x: {x.shape} | xi: {xi.shape}\")\n",
@@ -372,34 +403,35 @@
     "    return make_interp_spline(x, data, k=1, axis=-1)(xi)\n",
     "\n",
     "\n",
-    "interped = xr.apply_ufunc(\n",
-    "    debug_interp_batched,\n",
+    "interped_batched = xr.apply_ufunc(\n",
+    "    debug_interp_batched,  # only the function and vectorize changed\n",
     "    newlat_t,\n",
     "    air.lat,\n",
     "    air,\n",
     "    input_core_dims=[[\"lat\"], [\"lat\"], [\"lat\"]],\n",
     "    output_core_dims=[[\"lat\"]],\n",
-    "    exclude_dims={\"lat\"},\n",
+    "    exclude_dims={\"lat\"},  # dimensions allowed to change size. Must be set!\n",
     "    vectorize=[\"time\"],  # only loop over time\n",
     ")\n",
-    "interped"
+    "xr.testing.assert_allclose(interped, interped_batched)\n",
+    "interped_batched"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "16",
+   "id": "18",
    "metadata": {
     "tags": []
    },
    "source": [
-    "Now `debug_interp_batched` is called only 4 times, once per `time`, and receives all 3 `lon` at once: `data` has shape `(lon, lat)`.\n",
+    "Now `debug_interp_batched` is called only 4 times, once per `time`, and receives all 3 `lon` at once: `data` has shape `(lon, lat)`. The result is the same.\n",
     "\n",
     "Loop dimensions that are passed in bulk follow numpy's broadcasting rules. `xi` does not vary along `lon`, so it has size 1 along that dimension, and our function has to handle that."
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "17",
+   "id": "19",
    "metadata": {
     "tags": []
    },