{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "84bf8d8e-0536-4673-82be-55b6897627d6",
   "metadata": {},
   "source": [
    "# Dask support\n",
    "Like xarray, xarray-einstats also aims to support using their functions on\n",
    "NumPy or Dask backed xarray objects.\n",
    "\n",
    "Functions in xarray-einstats should support Dask backed xarray objects.\n",
    "But Dask support is still not extensively tested.\n",
    "If you have issues using Dask backed xarray objects please open an [issue](https://github.com/arviz-devs/xarray-einstats)\n",
    "\n",
    "Dask support within xarray-einstats requires explicit activation by users, exactly like it happens with {func}`xarray.apply_ufunc`.\n",
    "In fact, what is happening is xarray-einstats forwards any unrecognized keyword argument (e.g. `dask` or `dask_gufunc_kwargs`) to `apply_ufunc`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b373d879-c129-4441-9186-d4f98ddcb699",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dask.array as da\n",
    "import xarray as xr\n",
    "from xarray_einstats import linalg\n",
    "\n",
    "xr.set_options(display_expand_data=False);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e3daba12-78f5-473f-b9fc-71665434ec34",
   "metadata": {},
   "outputs": [],
   "source": [
    "matrix_batch = xr.DataArray(da.random.normal(size=(1000, 77, 77), chunks=\"auto\"), dims=[\"batch\", \"dim1\", \"dim2\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "724165af-0437-423f-8c22-3c0f5980eb52",
   "metadata": {},
   "source": [
    "Trying to use xarray-einstats functions directly on Dask arrays triggers an error:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "02bdc253-ff87-4481-979a-1db478f093d4",
   "metadata": {
    "deletable": true,
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "apply_ufunc encountered a chunked array on an argument, but handling for chunked arrays has not been enabled. Either set the ``dask`` argument or load your data into memory first with ``.load()`` or ``.compute()``",
     "output_type": "error",
     "traceback": [
      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
      "\u001b[31mValueError\u001b[39m                                Traceback (most recent call last)",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[3]\u001b[39m\u001b[32m, line 1\u001b[39m\n\u001b[32m----> \u001b[39m\u001b[32m1\u001b[39m linalg.norm(matrix_batch, dims=[\u001b[33m\"dim1\"\u001b[39m, \u001b[33m\"dim2\"\u001b[39m])\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/Documents/repos_oss/xarray-einstats/src/xarray_einstats/linalg.py:801\u001b[39m, in \u001b[36mnorm\u001b[39m\u001b[34m(da, dims, ord, **kwargs)\u001b[39m\n\u001b[32m    799\u001b[39m     in_dims = dims\n\u001b[32m    800\u001b[39m     norm_kwargs[\u001b[33m\"\u001b[39m\u001b[33maxis\u001b[39m\u001b[33m\"\u001b[39m] = (-\u001b[32m2\u001b[39m, -\u001b[32m1\u001b[39m)\n\u001b[32m--> \u001b[39m\u001b[32m801\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[30;43mxr\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mapply_ufunc\u001b[39;49m\u001b[30;43m(\u001b[39;49m\n\u001b[32m    802\u001b[39m \u001b[30;43m    \u001b[39;49m\u001b[30;43mnp\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mlinalg\u001b[39;49m\u001b[30;43m.\u001b[39;49m\u001b[30;43mnorm\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mda\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43minput_core_dims\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43m[\u001b[39;49m\u001b[30;43min_dims\u001b[39;49m\u001b[30;43m]\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43mkwargs\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43mnorm_kwargs\u001b[39;49m\u001b[30;43m,\u001b[39;49m\u001b[30;43m \u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43mkwargs\u001b[39;49m\n\u001b[32m    803\u001b[39m \u001b[30;43m\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/bin/miniforge3/envs/general/lib/python3.14/site-packages/xarray/computation/apply_ufunc.py:1267\u001b[39m, in \u001b[36mapply_ufunc\u001b[39m\u001b[34m(func, input_core_dims, output_core_dims, exclude_dims, vectorize, join, dataset_join, dataset_fill_value, keep_attrs, kwargs, dask, output_dtypes, output_sizes, meta, dask_gufunc_kwargs, on_missing_core_dim, *args)\u001b[39m\n\u001b[32m   1265\u001b[39m \u001b[38;5;66;03m# feed DataArray apply_variable_ufunc through apply_dataarray_vfunc\u001b[39;00m\n\u001b[32m   1266\u001b[39m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28many\u001b[39m(\u001b[38;5;28misinstance\u001b[39m(a, DataArray) \u001b[38;5;28;01mfor\u001b[39;00m a \u001b[38;5;129;01min\u001b[39;00m args):\n\u001b[32m-> \u001b[39m\u001b[32m1267\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[30;43mapply_dataarray_vfunc\u001b[39;49m\u001b[30;43m(\u001b[39;49m\n\u001b[32m   1268\u001b[39m \u001b[30;43m        \u001b[39;49m\u001b[30;43mvariables_vfunc\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m   1269\u001b[39m \u001b[30;43m        \u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43margs\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m   1270\u001b[39m \u001b[30;43m        \u001b[39;49m\u001b[30;43msignature\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43msignature\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m   1271\u001b[39m \u001b[30;43m        \u001b[39;49m\u001b[30;43mjoin\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43mjoin\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m   1272\u001b[39m \u001b[30;43m        \u001b[39;49m\u001b[30;43mexclude_dims\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43mexclude_dims\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m   1273\u001b[39m \u001b[30;43m        \u001b[39;49m\u001b[30;43mkeep_attrs\u001b[39;49m\u001b[30;43m=\u001b[39;49m\u001b[30;43mkeep_attrs\u001b[39;49m\u001b[30;43m,\u001b[39;49m\n\u001b[32m   1274\u001b[39m \u001b[30;43m    \u001b[39;49m\u001b[30;43m)\u001b[39;49m\n\u001b[32m   1275\u001b[39m \u001b[38;5;66;03m# feed Variables directly through apply_variable_ufunc\u001b[39;00m\n\u001b[32m   1276\u001b[39m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28many\u001b[39m(\u001b[38;5;28misinstance\u001b[39m(a, Variable) \u001b[38;5;28;01mfor\u001b[39;00m a \u001b[38;5;129;01min\u001b[39;00m args):\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/bin/miniforge3/envs/general/lib/python3.14/site-packages/xarray/computation/apply_ufunc.py:312\u001b[39m, in \u001b[36mapply_dataarray_vfunc\u001b[39m\u001b[34m(func, signature, join, exclude_dims, keep_attrs, *args)\u001b[39m\n\u001b[32m    307\u001b[39m result_coords, result_indexes = build_output_coords_and_indexes(\n\u001b[32m    308\u001b[39m     args, signature, exclude_dims, combine_attrs=keep_attrs\n\u001b[32m    309\u001b[39m )\n\u001b[32m    311\u001b[39m data_vars = [\u001b[38;5;28mgetattr\u001b[39m(a, \u001b[33m\"\u001b[39m\u001b[33mvariable\u001b[39m\u001b[33m\"\u001b[39m, a) \u001b[38;5;28;01mfor\u001b[39;00m a \u001b[38;5;129;01min\u001b[39;00m args]\n\u001b[32m--> \u001b[39m\u001b[32m312\u001b[39m result_var = \u001b[30;43mfunc\u001b[39;49m\u001b[30;43m(\u001b[39;49m\u001b[30;43m*\u001b[39;49m\u001b[30;43mdata_vars\u001b[39;49m\u001b[30;43m)\u001b[39;49m\n\u001b[32m    314\u001b[39m out: \u001b[38;5;28mtuple\u001b[39m[DataArray, ...] | DataArray\n\u001b[32m    315\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m signature.num_outputs > \u001b[32m1\u001b[39m:\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/bin/miniforge3/envs/general/lib/python3.14/site-packages/xarray/computation/apply_ufunc.py:741\u001b[39m, in \u001b[36mapply_variable_ufunc\u001b[39m\u001b[34m(func, signature, exclude_dims, dask, output_dtypes, vectorize, keep_attrs, dask_gufunc_kwargs, *args)\u001b[39m\n\u001b[32m    739\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28many\u001b[39m(is_chunked_array(array) \u001b[38;5;28;01mfor\u001b[39;00m array \u001b[38;5;129;01min\u001b[39;00m input_data):\n\u001b[32m    740\u001b[39m     \u001b[38;5;28;01mif\u001b[39;00m dask == \u001b[33m\"\u001b[39m\u001b[33mforbidden\u001b[39m\u001b[33m\"\u001b[39m:\n\u001b[32m--> \u001b[39m\u001b[32m741\u001b[39m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\n\u001b[32m    742\u001b[39m             \u001b[33m\"\u001b[39m\u001b[33mapply_ufunc encountered a chunked array on an \u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m    743\u001b[39m             \u001b[33m\"\u001b[39m\u001b[33margument, but handling for chunked arrays has not \u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m    744\u001b[39m             \u001b[33m\"\u001b[39m\u001b[33mbeen enabled. Either set the ``dask`` argument \u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m    745\u001b[39m             \u001b[33m\"\u001b[39m\u001b[33mor load your data into memory first with \u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m    746\u001b[39m             \u001b[33m\"\u001b[39m\u001b[33m``.load()`` or ``.compute()``\u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m    747\u001b[39m         )\n\u001b[32m    748\u001b[39m     \u001b[38;5;28;01melif\u001b[39;00m dask == \u001b[33m\"\u001b[39m\u001b[33mparallelized\u001b[39m\u001b[33m\"\u001b[39m:\n\u001b[32m    749\u001b[39m         chunkmanager = get_chunked_array_type(*input_data)\n",
      "\u001b[31mValueError\u001b[39m: apply_ufunc encountered a chunked array on an argument, but handling for chunked arrays has not been enabled. Either set the ``dask`` argument or load your data into memory first with ``.load()`` or ``.compute()``"
     ]
    }
   ],
   "source": [
    "linalg.norm(matrix_batch, dims=[\"dim1\", \"dim2\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "298e23b9-d271-4e68-a2c2-02dd661eb8bc",
   "metadata": {
    "deletable": true,
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "We need to allow dask inputs through `dask=\"allowed\"` or `dask=\"parallelized\"`. In general, `allowed` is preferred, but it only works if there is a Dask implementation of the requested function available, `parallelized` works in most cases but is usually slower and still not a silver bullet. Consequently, xarray-einstats does not set the `dask` argument automatically when the input is a dask array, and finding which of the two will work and be better might require a bit of trial an error an also be dependent on the Dask version available.\n",
    "\n",
    "Here is one example with `linalg.norm` for which both alternatives work:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b03c96a5-23b7-4fde-8ee9-b56911559a19",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table>\n",
       "    <tr>\n",
       "        <td>\n",
       "            <table style=\"border-collapse: collapse;\">\n",
       "                <thead>\n",
       "                    <tr>\n",
       "                        <td> </td>\n",
       "                        <th> Array </th>\n",
       "                        <th> Chunk </th>\n",
       "                    </tr>\n",
       "                </thead>\n",
       "                <tbody>\n",
       "                    \n",
       "                    <tr>\n",
       "                        <th> Bytes </th>\n",
       "                        <td> 7.81 kiB </td>\n",
       "                        <td> 7.81 kiB </td>\n",
       "                    </tr>\n",
       "                    \n",
       "                    <tr>\n",
       "                        <th> Shape </th>\n",
       "                        <td> (1000,) </td>\n",
       "                        <td> (1000,) </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <th> Dask graph </th>\n",
       "                        <td colspan=\"2\"> 1 chunks in 6 graph layers </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <th> Data type </th>\n",
       "                        <td colspan=\"2\"> float64 numpy.ndarray </td>\n",
       "                    </tr>\n",
       "                </tbody>\n",
       "            </table>\n",
       "        </td>\n",
       "        <td>\n",
       "        <svg width=\"170\" height=\"75\" style=\"stroke:rgb(0,0,0);stroke-width:1\" >\n",
       "\n",
       "  <!-- Horizontal lines -->\n",
       "  <line x1=\"0\" y1=\"0\" x2=\"120\" y2=\"0\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"0\" y1=\"25\" x2=\"120\" y2=\"25\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Vertical lines -->\n",
       "  <line x1=\"0\" y1=\"0\" x2=\"0\" y2=\"25\" style=\"stroke-width:2\" />\n",
       "  <line x1=\"120\" y1=\"0\" x2=\"120\" y2=\"25\" style=\"stroke-width:2\" />\n",
       "\n",
       "  <!-- Colored Rectangle -->\n",
       "  <polygon points=\"0.0,0.0 120.0,0.0 120.0,25.412616514582485 0.0,25.412616514582485\" style=\"fill:#ECB172A0;stroke-width:0\"/>\n",
       "\n",
       "  <!-- Text -->\n",
       "  <text x=\"60.0\" y=\"45.41261651458248\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" >1000</text>\n",
       "  <text x=\"140.0\" y=\"12.706308257291242\" font-size=\"1.0rem\" font-weight=\"100\" text-anchor=\"middle\" transform=\"rotate(0,140.0,12.706308257291242)\">1</text>\n",
       "</svg>\n",
       "        </td>\n",
       "    </tr>\n",
       "</table>"
      ],
      "text/plain": [
       "dask.array<pow, shape=(1000,), dtype=float64, chunksize=(1000,), chunktype=numpy.ndarray>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "norms = linalg.norm(matrix_batch, dims=[\"dim1\", \"dim2\"], dask=\"allowed\")\n",
    "norms.data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3287e425-a7da-4756-87d9-1b9b43e7bad5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><svg style=\"position: absolute; width: 0; height: 0; overflow: hidden\">\n",
       "<defs>\n",
       "<symbol id=\"icon-database\" viewBox=\"0 0 32 32\">\n",
       "<path d=\"M16 0c-8.837 0-16 2.239-16 5v4c0 2.761 7.163 5 16 5s16-2.239 16-5v-4c0-2.761-7.163-5-16-5z\"></path>\n",
       "<path d=\"M16 17c-8.837 0-16-2.239-16-5v6c0 2.761 7.163 5 16 5s16-2.239 16-5v-6c0 2.761-7.163 5-16 5z\"></path>\n",
       "<path d=\"M16 26c-8.837 0-16-2.239-16-5v6c0 2.761 7.163 5 16 5s16-2.239 16-5v-6c0 2.761-7.163 5-16 5z\"></path>\n",
       "</symbol>\n",
       "<symbol id=\"icon-file-text2\" viewBox=\"0 0 32 32\">\n",
       "<path d=\"M28.681 7.159c-0.694-0.947-1.662-2.053-2.724-3.116s-2.169-2.030-3.116-2.724c-1.612-1.182-2.393-1.319-2.841-1.319h-15.5c-1.378 0-2.5 1.121-2.5 2.5v27c0 1.378 1.122 2.5 2.5 2.5h23c1.378 0 2.5-1.122 2.5-2.5v-19.5c0-0.448-0.137-1.23-1.319-2.841zM24.543 5.457c0.959 0.959 1.712 1.825 2.268 2.543h-4.811v-4.811c0.718 0.556 1.584 1.309 2.543 2.268zM28 29.5c0 0.271-0.229 0.5-0.5 0.5h-23c-0.271 0-0.5-0.229-0.5-0.5v-27c0-0.271 0.229-0.5 0.5-0.5 0 0 15.499-0 15.5 0v7c0 0.552 0.448 1 1 1h7v19.5z\"></path>\n",
       "<path d=\"M23 26h-14c-0.552 0-1-0.448-1-1s0.448-1 1-1h14c0.552 0 1 0.448 1 1s-0.448 1-1 1z\"></path>\n",
       "<path d=\"M23 22h-14c-0.552 0-1-0.448-1-1s0.448-1 1-1h14c0.552 0 1 0.448 1 1s-0.448 1-1 1z\"></path>\n",
       "<path d=\"M23 18h-14c-0.552 0-1-0.448-1-1s0.448-1 1-1h14c0.552 0 1 0.448 1 1s-0.448 1-1 1z\"></path>\n",
       "</symbol>\n",
       "</defs>\n",
       "</svg>\n",
       "<style>/* CSS stylesheet for displaying xarray objects in notebooks */\n",
       "\n",
       ":root {\n",
       "  --xr-font-color0: var(\n",
       "    --jp-content-font-color0,\n",
       "    var(--pst-color-text-base rgba(0, 0, 0, 1))\n",
       "  );\n",
       "  --xr-font-color2: var(\n",
       "    --jp-content-font-color2,\n",
       "    var(--pst-color-text-base, rgba(0, 0, 0, 0.54))\n",
       "  );\n",
       "  --xr-font-color3: var(\n",
       "    --jp-content-font-color3,\n",
       "    var(--pst-color-text-base, rgba(0, 0, 0, 0.38))\n",
       "  );\n",
       "  --xr-border-color: var(\n",
       "    --jp-border-color2,\n",
       "    hsl(from var(--pst-color-on-background, white) h s calc(l - 10))\n",
       "  );\n",
       "  --xr-disabled-color: var(\n",
       "    --jp-layout-color3,\n",
       "    hsl(from var(--pst-color-on-background, white) h s calc(l - 40))\n",
       "  );\n",
       "  --xr-background-color: var(\n",
       "    --jp-layout-color0,\n",
       "    var(--pst-color-on-background, white)\n",
       "  );\n",
       "  --xr-background-color-row-even: var(\n",
       "    --jp-layout-color1,\n",
       "    hsl(from var(--pst-color-on-background, white) h s calc(l - 5))\n",
       "  );\n",
       "  --xr-background-color-row-odd: var(\n",
       "    --jp-layout-color2,\n",
       "    hsl(from var(--pst-color-on-background, white) h s calc(l - 15))\n",
       "  );\n",
       "}\n",
       "\n",
       "html[theme=\"dark\"],\n",
       "html[data-theme=\"dark\"],\n",
       "body[data-theme=\"dark\"],\n",
       "body.vscode-dark {\n",
       "  --xr-font-color0: var(\n",
       "    --jp-content-font-color0,\n",
       "    var(--pst-color-text-base, rgba(255, 255, 255, 1))\n",
       "  );\n",
       "  --xr-font-color2: var(\n",
       "    --jp-content-font-color2,\n",
       "    var(--pst-color-text-base, rgba(255, 255, 255, 0.54))\n",
       "  );\n",
       "  --xr-font-color3: var(\n",
       "    --jp-content-font-color3,\n",
       "    var(--pst-color-text-base, rgba(255, 255, 255, 0.38))\n",
       "  );\n",
       "  --xr-border-color: var(\n",
       "    --jp-border-color2,\n",
       "    hsl(from var(--pst-color-on-background, #111111) h s calc(l + 10))\n",
       "  );\n",
       "  --xr-disabled-color: var(\n",
       "    --jp-layout-color3,\n",
       "    hsl(from var(--pst-color-on-background, #111111) h s calc(l + 40))\n",
       "  );\n",
       "  --xr-background-color: var(\n",
       "    --jp-layout-color0,\n",
       "    var(--pst-color-on-background, #111111)\n",
       "  );\n",
       "  --xr-background-color-row-even: var(\n",
       "    --jp-layout-color1,\n",
       "    hsl(from var(--pst-color-on-background, #111111) h s calc(l + 5))\n",
       "  );\n",
       "  --xr-background-color-row-odd: var(\n",
       "    --jp-layout-color2,\n",
       "    hsl(from var(--pst-color-on-background, #111111) h s calc(l + 15))\n",
       "  );\n",
       "}\n",
       "\n",
       ".xr-wrap {\n",
       "  display: block !important;\n",
       "  min-width: 300px;\n",
       "  max-width: 700px;\n",
       "  line-height: 1.6;\n",
       "  padding-bottom: 4px;\n",
       "}\n",
       "\n",
       ".xr-text-repr-fallback {\n",
       "  /* fallback to plain text repr when CSS is not injected (untrusted notebook) */\n",
       "  display: none;\n",
       "}\n",
       "\n",
       ".xr-header {\n",
       "  padding-top: 6px;\n",
       "  padding-bottom: 6px;\n",
       "}\n",
       "\n",
       ".xr-header {\n",
       "  border-bottom: solid 1px var(--xr-border-color);\n",
       "  margin-bottom: 4px;\n",
       "}\n",
       "\n",
       ".xr-header > div,\n",
       ".xr-header > ul {\n",
       "  display: inline;\n",
       "  margin-top: 0;\n",
       "  margin-bottom: 0;\n",
       "}\n",
       "\n",
       ".xr-obj-type,\n",
       ".xr-obj-name {\n",
       "  margin-left: 2px;\n",
       "  margin-right: 10px;\n",
       "}\n",
       "\n",
       ".xr-obj-type,\n",
       ".xr-group-box-contents > label {\n",
       "  color: var(--xr-font-color2);\n",
       "  display: block;\n",
       "}\n",
       "\n",
       ".xr-sections {\n",
       "  padding-left: 0 !important;\n",
       "  display: grid;\n",
       "  grid-template-columns: 150px auto auto 1fr 0 20px 0 20px;\n",
       "  margin-block-start: 0;\n",
       "  margin-block-end: 0;\n",
       "}\n",
       "\n",
       ".xr-section-item {\n",
       "  display: contents;\n",
       "}\n",
       "\n",
       ".xr-section-item > input,\n",
       ".xr-group-box-contents > input,\n",
       ".xr-array-wrap > input {\n",
       "  display: block;\n",
       "  opacity: 0;\n",
       "  height: 0;\n",
       "  margin: 0;\n",
       "}\n",
       "\n",
       ".xr-section-item > input + label,\n",
       ".xr-var-item > input + label {\n",
       "  color: var(--xr-disabled-color);\n",
       "}\n",
       "\n",
       ".xr-section-item > input:enabled + label,\n",
       ".xr-var-item > input:enabled + label,\n",
       ".xr-array-wrap > input:enabled + label,\n",
       ".xr-group-box-contents > input:enabled + label {\n",
       "  cursor: pointer;\n",
       "  color: var(--xr-font-color2);\n",
       "}\n",
       "\n",
       ".xr-section-item > input:focus-visible + label,\n",
       ".xr-var-item > input:focus-visible + label,\n",
       ".xr-array-wrap > input:focus-visible + label,\n",
       ".xr-group-box-contents > input:focus-visible + label {\n",
       "  outline: auto;\n",
       "}\n",
       "\n",
       ".xr-section-item > input:enabled + label:hover,\n",
       ".xr-var-item > input:enabled + label:hover,\n",
       ".xr-array-wrap > input:enabled + label:hover,\n",
       ".xr-group-box-contents > input:enabled + label:hover {\n",
       "  color: var(--xr-font-color0);\n",
       "}\n",
       "\n",
       ".xr-section-summary {\n",
       "  grid-column: 1;\n",
       "  color: var(--xr-font-color2);\n",
       "  font-weight: 500;\n",
       "  white-space: nowrap;\n",
       "}\n",
       "\n",
       ".xr-section-summary > em {\n",
       "  font-weight: normal;\n",
       "}\n",
       "\n",
       ".xr-span-grid {\n",
       "  grid-column-end: -1;\n",
       "}\n",
       "\n",
       ".xr-section-summary > span {\n",
       "  display: inline-block;\n",
       "  padding-left: 0.3em;\n",
       "}\n",
       "\n",
       ".xr-group-box-contents > input:checked + label > span {\n",
       "  display: inline-block;\n",
       "  padding-left: 0.6em;\n",
       "}\n",
       "\n",
       ".xr-section-summary-in:disabled + label {\n",
       "  color: var(--xr-font-color2);\n",
       "}\n",
       "\n",
       ".xr-section-summary-in + label:before {\n",
       "  display: inline-block;\n",
       "  content: \"►\";\n",
       "  font-size: 11px;\n",
       "  width: 15px;\n",
       "  text-align: center;\n",
       "}\n",
       "\n",
       ".xr-section-summary-in:disabled + label:before {\n",
       "  color: var(--xr-disabled-color);\n",
       "}\n",
       "\n",
       ".xr-section-summary-in:checked + label:before {\n",
       "  content: \"▼\";\n",
       "}\n",
       "\n",
       ".xr-section-summary-in:checked + label > span {\n",
       "  display: none;\n",
       "}\n",
       "\n",
       ".xr-section-summary,\n",
       ".xr-section-inline-details,\n",
       ".xr-group-box-contents > label {\n",
       "  padding-top: 4px;\n",
       "}\n",
       "\n",
       ".xr-section-inline-details {\n",
       "  grid-column: 2 / -1;\n",
       "}\n",
       "\n",
       ".xr-section-details {\n",
       "  grid-column: 1 / -1;\n",
       "  margin-top: 4px;\n",
       "  margin-bottom: 5px;\n",
       "}\n",
       "\n",
       ".xr-section-summary-in ~ .xr-section-details {\n",
       "  display: none;\n",
       "}\n",
       "\n",
       ".xr-section-summary-in:checked ~ .xr-section-details {\n",
       "  display: contents;\n",
       "}\n",
       "\n",
       ".xr-children {\n",
       "  display: inline-grid;\n",
       "  grid-template-columns: 100%;\n",
       "  grid-column: 1 / -1;\n",
       "  padding-top: 4px;\n",
       "}\n",
       "\n",
       ".xr-group-box {\n",
       "  display: inline-grid;\n",
       "  grid-template-columns: 0px 30px auto;\n",
       "}\n",
       "\n",
       ".xr-group-box-vline {\n",
       "  grid-column-start: 1;\n",
       "  border-right: 0.2em solid;\n",
       "  border-color: var(--xr-border-color);\n",
       "  width: 0px;\n",
       "}\n",
       "\n",
       ".xr-group-box-hline {\n",
       "  grid-column-start: 2;\n",
       "  grid-row-start: 1;\n",
       "  height: 1em;\n",
       "  width: 26px;\n",
       "  border-bottom: 0.2em solid;\n",
       "  border-color: var(--xr-border-color);\n",
       "}\n",
       "\n",
       ".xr-group-box-contents {\n",
       "  grid-column-start: 3;\n",
       "  padding-bottom: 4px;\n",
       "}\n",
       "\n",
       ".xr-group-box-contents > label::before {\n",
       "  content: \"📂\";\n",
       "  padding-right: 0.3em;\n",
       "}\n",
       "\n",
       ".xr-group-box-contents > input:checked + label::before {\n",
       "  content: \"📁\";\n",
       "}\n",
       "\n",
       ".xr-group-box-contents > input:checked + label {\n",
       "  padding-bottom: 0px;\n",
       "}\n",
       "\n",
       ".xr-group-box-contents > input:checked ~ .xr-sections {\n",
       "  display: none;\n",
       "}\n",
       "\n",
       ".xr-group-box-contents > input + label > span {\n",
       "  display: none;\n",
       "}\n",
       "\n",
       ".xr-group-box-ellipsis {\n",
       "  font-size: 1.4em;\n",
       "  font-weight: 900;\n",
       "  color: var(--xr-font-color2);\n",
       "  letter-spacing: 0.15em;\n",
       "  cursor: default;\n",
       "}\n",
       "\n",
       ".xr-array-wrap {\n",
       "  grid-column: 1 / -1;\n",
       "  display: grid;\n",
       "  grid-template-columns: 20px auto;\n",
       "}\n",
       "\n",
       ".xr-array-wrap > label {\n",
       "  grid-column: 1;\n",
       "  vertical-align: top;\n",
       "}\n",
       "\n",
       ".xr-preview {\n",
       "  color: var(--xr-font-color3);\n",
       "}\n",
       "\n",
       ".xr-array-preview,\n",
       ".xr-array-data {\n",
       "  padding: 0 5px !important;\n",
       "  grid-column: 2;\n",
       "}\n",
       "\n",
       ".xr-array-data,\n",
       ".xr-array-in:checked ~ .xr-array-preview {\n",
       "  display: none;\n",
       "}\n",
       "\n",
       ".xr-array-in:checked ~ .xr-array-data,\n",
       ".xr-array-preview {\n",
       "  display: inline-block;\n",
       "}\n",
       "\n",
       ".xr-dim-list {\n",
       "  display: inline-block !important;\n",
       "  list-style: none;\n",
       "  padding: 0 !important;\n",
       "  margin: 0;\n",
       "}\n",
       "\n",
       ".xr-dim-list li {\n",
       "  display: inline-block;\n",
       "  padding: 0;\n",
       "  margin: 0;\n",
       "}\n",
       "\n",
       ".xr-dim-list:before {\n",
       "  content: \"(\";\n",
       "}\n",
       "\n",
       ".xr-dim-list:after {\n",
       "  content: \")\";\n",
       "}\n",
       "\n",
       ".xr-dim-list li:not(:last-child):after {\n",
       "  content: \",\";\n",
       "  padding-right: 5px;\n",
       "}\n",
       "\n",
       ".xr-has-index {\n",
       "  font-weight: bold;\n",
       "}\n",
       "\n",
       ".xr-var-list,\n",
       ".xr-var-item {\n",
       "  display: contents;\n",
       "}\n",
       "\n",
       ".xr-var-item > div,\n",
       ".xr-var-item label,\n",
       ".xr-var-item > .xr-var-name span {\n",
       "  background-color: var(--xr-background-color-row-even);\n",
       "  border-color: var(--xr-background-color-row-odd);\n",
       "  margin-bottom: 0;\n",
       "  padding-top: 2px;\n",
       "}\n",
       "\n",
       ".xr-var-item > .xr-var-name:hover span {\n",
       "  padding-right: 5px;\n",
       "}\n",
       "\n",
       ".xr-var-list > li:nth-child(odd) > div,\n",
       ".xr-var-list > li:nth-child(odd) > label,\n",
       ".xr-var-list > li:nth-child(odd) > .xr-var-name span {\n",
       "  background-color: var(--xr-background-color-row-odd);\n",
       "  border-color: var(--xr-background-color-row-even);\n",
       "}\n",
       "\n",
       ".xr-var-name {\n",
       "  grid-column: 1;\n",
       "}\n",
       "\n",
       ".xr-var-dims {\n",
       "  grid-column: 2;\n",
       "}\n",
       "\n",
       ".xr-var-dtype {\n",
       "  grid-column: 3;\n",
       "  text-align: right;\n",
       "  color: var(--xr-font-color2);\n",
       "}\n",
       "\n",
       ".xr-var-preview {\n",
       "  grid-column: 4;\n",
       "}\n",
       "\n",
       ".xr-index-preview {\n",
       "  grid-column: 2 / 5;\n",
       "  color: var(--xr-font-color2);\n",
       "}\n",
       "\n",
       ".xr-var-name,\n",
       ".xr-var-dims,\n",
       ".xr-var-dtype,\n",
       ".xr-preview,\n",
       ".xr-attrs dt {\n",
       "  white-space: nowrap;\n",
       "  overflow: hidden;\n",
       "  text-overflow: ellipsis;\n",
       "  padding-right: 10px;\n",
       "}\n",
       "\n",
       ".xr-var-name:hover,\n",
       ".xr-var-dims:hover,\n",
       ".xr-var-dtype:hover,\n",
       ".xr-attrs dt:hover {\n",
       "  overflow: visible;\n",
       "  width: auto;\n",
       "  z-index: 1;\n",
       "}\n",
       "\n",
       ".xr-var-attrs,\n",
       ".xr-var-data,\n",
       ".xr-index-data {\n",
       "  display: none;\n",
       "  border-top: 2px dotted var(--xr-background-color);\n",
       "  padding-bottom: 20px !important;\n",
       "  padding-top: 10px !important;\n",
       "}\n",
       "\n",
       ".xr-var-attrs-in + label,\n",
       ".xr-var-data-in + label,\n",
       ".xr-index-data-in + label {\n",
       "  padding: 0 1px;\n",
       "}\n",
       "\n",
       ".xr-var-attrs-in:checked ~ .xr-var-attrs,\n",
       ".xr-var-data-in:checked ~ .xr-var-data,\n",
       ".xr-index-data-in:checked ~ .xr-index-data {\n",
       "  display: block;\n",
       "}\n",
       "\n",
       ".xr-var-data > table {\n",
       "  float: right;\n",
       "}\n",
       "\n",
       ".xr-var-data > pre,\n",
       ".xr-index-data > pre,\n",
       ".xr-var-data > table > tbody > tr {\n",
       "  background-color: transparent !important;\n",
       "}\n",
       "\n",
       ".xr-var-name span,\n",
       ".xr-var-data,\n",
       ".xr-index-name div,\n",
       ".xr-index-data,\n",
       ".xr-attrs {\n",
       "  padding-left: 25px !important;\n",
       "}\n",
       "\n",
       ".xr-attrs,\n",
       ".xr-var-attrs,\n",
       ".xr-var-data,\n",
       ".xr-index-data {\n",
       "  grid-column: 1 / -1;\n",
       "}\n",
       "\n",
       "dl.xr-attrs {\n",
       "  padding: 0;\n",
       "  margin: 0;\n",
       "  display: grid;\n",
       "  grid-template-columns: 125px auto;\n",
       "}\n",
       "\n",
       ".xr-attrs dt,\n",
       ".xr-attrs dd {\n",
       "  padding: 0;\n",
       "  margin: 0;\n",
       "  float: left;\n",
       "  padding-right: 10px;\n",
       "  width: auto;\n",
       "}\n",
       "\n",
       ".xr-attrs dt {\n",
       "  font-weight: normal;\n",
       "  grid-column: 1;\n",
       "}\n",
       "\n",
       ".xr-attrs dt:hover span {\n",
       "  display: inline-block;\n",
       "  background: var(--xr-background-color);\n",
       "  padding-right: 10px;\n",
       "}\n",
       "\n",
       ".xr-attrs dd {\n",
       "  grid-column: 2;\n",
       "  white-space: pre-wrap;\n",
       "  word-break: break-all;\n",
       "}\n",
       "\n",
       ".xr-icon-database,\n",
       ".xr-icon-file-text2,\n",
       ".xr-no-icon {\n",
       "  display: inline-block;\n",
       "  vertical-align: middle;\n",
       "  width: 1em;\n",
       "  height: 1.5em !important;\n",
       "  stroke-width: 0;\n",
       "  stroke: currentColor;\n",
       "  fill: currentColor;\n",
       "}\n",
       "\n",
       ".xr-var-attrs-in:checked + label > .xr-icon-file-text2,\n",
       ".xr-var-data-in:checked + label > .xr-icon-database,\n",
       ".xr-index-data-in:checked + label > .xr-icon-database {\n",
       "  color: var(--xr-font-color0);\n",
       "  filter: drop-shadow(1px 1px 5px var(--xr-font-color2));\n",
       "  stroke-width: 0.8px;\n",
       "}\n",
       "</style><pre class='xr-text-repr-fallback'>&lt;xarray.DataArray &#x27;normal-758a57bb75ed3127c94212c3e83a5ed8&#x27; (batch: 1000)&gt; Size: 8kB\n",
       "77.07 77.61 76.86 77.53 77.99 77.22 76.9 ... 77.92 78.03 77.96 76.6 77.31 76.26\n",
       "Dimensions without coordinates: batch</pre><div class='xr-wrap' style='display:none'><div class='xr-header'><div class='xr-obj-type'>xarray.DataArray</div><div class='xr-obj-name'>&#x27;normal-758a57bb75ed3127c94212c3e83a5ed8&#x27;</div><ul class='xr-dim-list'><li><span>batch</span>: 1000</li></ul></div><ul class='xr-sections'><li class='xr-section-item'><div class='xr-array-wrap'><input id='section-df0196f0-371a-4681-8224-2eb46aae1573' class='xr-array-in' type='checkbox' ><label for='section-df0196f0-371a-4681-8224-2eb46aae1573' title='Show/hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-array-preview xr-preview'><span>77.07 77.61 76.86 77.53 77.99 77.22 ... 78.03 77.96 76.6 77.31 76.26</span></div><div class='xr-array-data'><pre>array([77.07151123, 77.61059979, 76.86360161, 77.52666541, 77.98533059,\n",
       "       77.22236039, 76.89694279, 76.5113661 , 77.04069162, 77.0440349 ,\n",
       "       76.38611509, 77.30566883, 77.47694384, 76.69118445, 76.3648431 ,\n",
       "       78.65136421, 75.48230116, 77.29719159, 76.69989052, 77.53014796,\n",
       "       78.12950831, 76.58109187, 76.1735932 , 77.57211317, 76.31369742,\n",
       "       76.28008721, 76.00930925, 77.69946061, 78.47248892, 76.30594417,\n",
       "       76.74855535, 77.23211165, 77.09690595, 77.58629277, 77.25236896,\n",
       "       77.29720169, 77.70320198, 77.29337301, 77.51158852, 77.00113117,\n",
       "       77.61088289, 76.56355944, 76.66462471, 78.35684907, 77.62398703,\n",
       "       77.15455699, 77.42167832, 76.20287995, 75.56872634, 77.17161862,\n",
       "       77.6550905 , 77.3271054 , 77.00974287, 78.52692208, 76.62211391,\n",
       "       76.63936546, 77.72529412, 75.9304276 , 77.94115594, 76.28197731,\n",
       "       77.86032362, 77.34556766, 78.28347133, 78.25610074, 76.4990315 ,\n",
       "       77.05555913, 77.7762863 , 76.83277802, 75.61555854, 77.93083141,\n",
       "       76.35444611, 76.40108947, 77.10925417, 76.21775021, 76.50717149,\n",
       "       78.62424972, 78.41599829, 77.53805218, 76.29311862, 77.23899266,\n",
       "       77.17823025, 76.23308925, 77.10158854, 78.06501202, 77.81267966,\n",
       "       75.9464039 , 77.37662775, 76.22437435, 76.93255673, 77.05499553,\n",
       "       78.42548546, 76.63240287, 78.43203628, 77.27066793, 76.95095667,\n",
       "       77.6427735 , 78.23776876, 77.98181894, 76.55946822, 76.59540871,\n",
       "...\n",
       "       77.55074864, 77.08245593, 76.05782067, 77.52499048, 78.07990445,\n",
       "       76.95812276, 76.09941659, 77.63809857, 75.79134748, 77.27308766,\n",
       "       77.26043898, 76.87723759, 76.20732289, 76.91405135, 78.67692647,\n",
       "       76.95501719, 77.22709319, 76.90604196, 77.1620702 , 75.95007881,\n",
       "       76.99028891, 77.16984169, 77.32234885, 76.23063511, 76.81834779,\n",
       "       76.15704279, 75.61472124, 78.17831097, 76.7443434 , 76.50639742,\n",
       "       77.62027759, 75.975392  , 76.57688827, 77.47909537, 77.36364272,\n",
       "       77.42555982, 76.23375548, 76.81310853, 77.87934912, 76.94914194,\n",
       "       77.22891481, 77.04548581, 77.41784027, 76.99567393, 75.91612618,\n",
       "       77.92538295, 77.87590131, 75.91280074, 77.01375182, 76.84293853,\n",
       "       77.29023139, 75.74463574, 76.52104818, 76.9560999 , 76.93174665,\n",
       "       75.94148142, 75.8301277 , 75.99251482, 77.34913468, 78.26571634,\n",
       "       77.86725889, 76.94211159, 76.92407092, 76.70451861, 77.12471401,\n",
       "       77.26577934, 76.06551508, 76.75742254, 76.21261178, 76.68219503,\n",
       "       76.59003924, 76.59220068, 75.72158433, 76.32204662, 77.53317453,\n",
       "       77.35568039, 78.39380204, 78.01003637, 76.57906144, 76.44788913,\n",
       "       76.16136169, 76.78692734, 77.15260584, 76.97373823, 78.16022121,\n",
       "       77.12782383, 76.94807175, 76.84872826, 76.92376664, 75.95237816,\n",
       "       76.72480888, 76.89486946, 77.19022507, 77.66279427, 77.91773289,\n",
       "       78.02838772, 77.96117143, 76.60281147, 77.3129846 , 76.25648097])</pre></div></div></li></ul></div></div>"
      ],
      "text/plain": [
       "<xarray.DataArray 'normal-758a57bb75ed3127c94212c3e83a5ed8' (batch: 1000)> Size: 8kB\n",
       "77.07 77.61 76.86 77.53 77.99 77.22 76.9 ... 77.92 78.03 77.96 76.6 77.31 76.26\n",
       "Dimensions without coordinates: batch"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "norms.compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b19b98df-3256-443f-b6e4-6501d76b6bc1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><svg style=\"position: absolute; width: 0; height: 0; overflow: hidden\">\n",
       "<defs>\n",
       "<symbol id=\"icon-database\" viewBox=\"0 0 32 32\">\n",
       "<path d=\"M16 0c-8.837 0-16 2.239-16 5v4c0 2.761 7.163 5 16 5s16-2.239 16-5v-4c0-2.761-7.163-5-16-5z\"></path>\n",
       "<path d=\"M16 17c-8.837 0-16-2.239-16-5v6c0 2.761 7.163 5 16 5s16-2.239 16-5v-6c0 2.761-7.163 5-16 5z\"></path>\n",
       "<path d=\"M16 26c-8.837 0-16-2.239-16-5v6c0 2.761 7.163 5 16 5s16-2.239 16-5v-6c0 2.761-7.163 5-16 5z\"></path>\n",
       "</symbol>\n",
       "<symbol id=\"icon-file-text2\" viewBox=\"0 0 32 32\">\n",
       "<path d=\"M28.681 7.159c-0.694-0.947-1.662-2.053-2.724-3.116s-2.169-2.030-3.116-2.724c-1.612-1.182-2.393-1.319-2.841-1.319h-15.5c-1.378 0-2.5 1.121-2.5 2.5v27c0 1.378 1.122 2.5 2.5 2.5h23c1.378 0 2.5-1.122 2.5-2.5v-19.5c0-0.448-0.137-1.23-1.319-2.841zM24.543 5.457c0.959 0.959 1.712 1.825 2.268 2.543h-4.811v-4.811c0.718 0.556 1.584 1.309 2.543 2.268zM28 29.5c0 0.271-0.229 0.5-0.5 0.5h-23c-0.271 0-0.5-0.229-0.5-0.5v-27c0-0.271 0.229-0.5 0.5-0.5 0 0 15.499-0 15.5 0v7c0 0.552 0.448 1 1 1h7v19.5z\"></path>\n",
       "<path d=\"M23 26h-14c-0.552 0-1-0.448-1-1s0.448-1 1-1h14c0.552 0 1 0.448 1 1s-0.448 1-1 1z\"></path>\n",
       "<path d=\"M23 22h-14c-0.552 0-1-0.448-1-1s0.448-1 1-1h14c0.552 0 1 0.448 1 1s-0.448 1-1 1z\"></path>\n",
       "<path d=\"M23 18h-14c-0.552 0-1-0.448-1-1s0.448-1 1-1h14c0.552 0 1 0.448 1 1s-0.448 1-1 1z\"></path>\n",
       "</symbol>\n",
       "</defs>\n",
       "</svg>\n",
       "<style>/* CSS stylesheet for displaying xarray objects in notebooks */\n",
       "\n",
       ":root {\n",
       "  --xr-font-color0: var(\n",
       "    --jp-content-font-color0,\n",
       "    var(--pst-color-text-base rgba(0, 0, 0, 1))\n",
       "  );\n",
       "  --xr-font-color2: var(\n",
       "    --jp-content-font-color2,\n",
       "    var(--pst-color-text-base, rgba(0, 0, 0, 0.54))\n",
       "  );\n",
       "  --xr-font-color3: var(\n",
       "    --jp-content-font-color3,\n",
       "    var(--pst-color-text-base, rgba(0, 0, 0, 0.38))\n",
       "  );\n",
       "  --xr-border-color: var(\n",
       "    --jp-border-color2,\n",
       "    hsl(from var(--pst-color-on-background, white) h s calc(l - 10))\n",
       "  );\n",
       "  --xr-disabled-color: var(\n",
       "    --jp-layout-color3,\n",
       "    hsl(from var(--pst-color-on-background, white) h s calc(l - 40))\n",
       "  );\n",
       "  --xr-background-color: var(\n",
       "    --jp-layout-color0,\n",
       "    var(--pst-color-on-background, white)\n",
       "  );\n",
       "  --xr-background-color-row-even: var(\n",
       "    --jp-layout-color1,\n",
       "    hsl(from var(--pst-color-on-background, white) h s calc(l - 5))\n",
       "  );\n",
       "  --xr-background-color-row-odd: var(\n",
       "    --jp-layout-color2,\n",
       "    hsl(from var(--pst-color-on-background, white) h s calc(l - 15))\n",
       "  );\n",
       "}\n",
       "\n",
       "html[theme=\"dark\"],\n",
       "html[data-theme=\"dark\"],\n",
       "body[data-theme=\"dark\"],\n",
       "body.vscode-dark {\n",
       "  --xr-font-color0: var(\n",
       "    --jp-content-font-color0,\n",
       "    var(--pst-color-text-base, rgba(255, 255, 255, 1))\n",
       "  );\n",
       "  --xr-font-color2: var(\n",
       "    --jp-content-font-color2,\n",
       "    var(--pst-color-text-base, rgba(255, 255, 255, 0.54))\n",
       "  );\n",
       "  --xr-font-color3: var(\n",
       "    --jp-content-font-color3,\n",
       "    var(--pst-color-text-base, rgba(255, 255, 255, 0.38))\n",
       "  );\n",
       "  --xr-border-color: var(\n",
       "    --jp-border-color2,\n",
       "    hsl(from var(--pst-color-on-background, #111111) h s calc(l + 10))\n",
       "  );\n",
       "  --xr-disabled-color: var(\n",
       "    --jp-layout-color3,\n",
       "    hsl(from var(--pst-color-on-background, #111111) h s calc(l + 40))\n",
       "  );\n",
       "  --xr-background-color: var(\n",
       "    --jp-layout-color0,\n",
       "    var(--pst-color-on-background, #111111)\n",
       "  );\n",
       "  --xr-background-color-row-even: var(\n",
       "    --jp-layout-color1,\n",
       "    hsl(from var(--pst-color-on-background, #111111) h s calc(l + 5))\n",
       "  );\n",
       "  --xr-background-color-row-odd: var(\n",
       "    --jp-layout-color2,\n",
       "    hsl(from var(--pst-color-on-background, #111111) h s calc(l + 15))\n",
       "  );\n",
       "}\n",
       "\n",
       ".xr-wrap {\n",
       "  display: block !important;\n",
       "  min-width: 300px;\n",
       "  max-width: 700px;\n",
       "  line-height: 1.6;\n",
       "  padding-bottom: 4px;\n",
       "}\n",
       "\n",
       ".xr-text-repr-fallback {\n",
       "  /* fallback to plain text repr when CSS is not injected (untrusted notebook) */\n",
       "  display: none;\n",
       "}\n",
       "\n",
       ".xr-header {\n",
       "  padding-top: 6px;\n",
       "  padding-bottom: 6px;\n",
       "}\n",
       "\n",
       ".xr-header {\n",
       "  border-bottom: solid 1px var(--xr-border-color);\n",
       "  margin-bottom: 4px;\n",
       "}\n",
       "\n",
       ".xr-header > div,\n",
       ".xr-header > ul {\n",
       "  display: inline;\n",
       "  margin-top: 0;\n",
       "  margin-bottom: 0;\n",
       "}\n",
       "\n",
       ".xr-obj-type,\n",
       ".xr-obj-name {\n",
       "  margin-left: 2px;\n",
       "  margin-right: 10px;\n",
       "}\n",
       "\n",
       ".xr-obj-type,\n",
       ".xr-group-box-contents > label {\n",
       "  color: var(--xr-font-color2);\n",
       "  display: block;\n",
       "}\n",
       "\n",
       ".xr-sections {\n",
       "  padding-left: 0 !important;\n",
       "  display: grid;\n",
       "  grid-template-columns: 150px auto auto 1fr 0 20px 0 20px;\n",
       "  margin-block-start: 0;\n",
       "  margin-block-end: 0;\n",
       "}\n",
       "\n",
       ".xr-section-item {\n",
       "  display: contents;\n",
       "}\n",
       "\n",
       ".xr-section-item > input,\n",
       ".xr-group-box-contents > input,\n",
       ".xr-array-wrap > input {\n",
       "  display: block;\n",
       "  opacity: 0;\n",
       "  height: 0;\n",
       "  margin: 0;\n",
       "}\n",
       "\n",
       ".xr-section-item > input + label,\n",
       ".xr-var-item > input + label {\n",
       "  color: var(--xr-disabled-color);\n",
       "}\n",
       "\n",
       ".xr-section-item > input:enabled + label,\n",
       ".xr-var-item > input:enabled + label,\n",
       ".xr-array-wrap > input:enabled + label,\n",
       ".xr-group-box-contents > input:enabled + label {\n",
       "  cursor: pointer;\n",
       "  color: var(--xr-font-color2);\n",
       "}\n",
       "\n",
       ".xr-section-item > input:focus-visible + label,\n",
       ".xr-var-item > input:focus-visible + label,\n",
       ".xr-array-wrap > input:focus-visible + label,\n",
       ".xr-group-box-contents > input:focus-visible + label {\n",
       "  outline: auto;\n",
       "}\n",
       "\n",
       ".xr-section-item > input:enabled + label:hover,\n",
       ".xr-var-item > input:enabled + label:hover,\n",
       ".xr-array-wrap > input:enabled + label:hover,\n",
       ".xr-group-box-contents > input:enabled + label:hover {\n",
       "  color: var(--xr-font-color0);\n",
       "}\n",
       "\n",
       ".xr-section-summary {\n",
       "  grid-column: 1;\n",
       "  color: var(--xr-font-color2);\n",
       "  font-weight: 500;\n",
       "  white-space: nowrap;\n",
       "}\n",
       "\n",
       ".xr-section-summary > em {\n",
       "  font-weight: normal;\n",
       "}\n",
       "\n",
       ".xr-span-grid {\n",
       "  grid-column-end: -1;\n",
       "}\n",
       "\n",
       ".xr-section-summary > span {\n",
       "  display: inline-block;\n",
       "  padding-left: 0.3em;\n",
       "}\n",
       "\n",
       ".xr-group-box-contents > input:checked + label > span {\n",
       "  display: inline-block;\n",
       "  padding-left: 0.6em;\n",
       "}\n",
       "\n",
       ".xr-section-summary-in:disabled + label {\n",
       "  color: var(--xr-font-color2);\n",
       "}\n",
       "\n",
       ".xr-section-summary-in + label:before {\n",
       "  display: inline-block;\n",
       "  content: \"►\";\n",
       "  font-size: 11px;\n",
       "  width: 15px;\n",
       "  text-align: center;\n",
       "}\n",
       "\n",
       ".xr-section-summary-in:disabled + label:before {\n",
       "  color: var(--xr-disabled-color);\n",
       "}\n",
       "\n",
       ".xr-section-summary-in:checked + label:before {\n",
       "  content: \"▼\";\n",
       "}\n",
       "\n",
       ".xr-section-summary-in:checked + label > span {\n",
       "  display: none;\n",
       "}\n",
       "\n",
       ".xr-section-summary,\n",
       ".xr-section-inline-details,\n",
       ".xr-group-box-contents > label {\n",
       "  padding-top: 4px;\n",
       "}\n",
       "\n",
       ".xr-section-inline-details {\n",
       "  grid-column: 2 / -1;\n",
       "}\n",
       "\n",
       ".xr-section-details {\n",
       "  grid-column: 1 / -1;\n",
       "  margin-top: 4px;\n",
       "  margin-bottom: 5px;\n",
       "}\n",
       "\n",
       ".xr-section-summary-in ~ .xr-section-details {\n",
       "  display: none;\n",
       "}\n",
       "\n",
       ".xr-section-summary-in:checked ~ .xr-section-details {\n",
       "  display: contents;\n",
       "}\n",
       "\n",
       ".xr-children {\n",
       "  display: inline-grid;\n",
       "  grid-template-columns: 100%;\n",
       "  grid-column: 1 / -1;\n",
       "  padding-top: 4px;\n",
       "}\n",
       "\n",
       ".xr-group-box {\n",
       "  display: inline-grid;\n",
       "  grid-template-columns: 0px 30px auto;\n",
       "}\n",
       "\n",
       ".xr-group-box-vline {\n",
       "  grid-column-start: 1;\n",
       "  border-right: 0.2em solid;\n",
       "  border-color: var(--xr-border-color);\n",
       "  width: 0px;\n",
       "}\n",
       "\n",
       ".xr-group-box-hline {\n",
       "  grid-column-start: 2;\n",
       "  grid-row-start: 1;\n",
       "  height: 1em;\n",
       "  width: 26px;\n",
       "  border-bottom: 0.2em solid;\n",
       "  border-color: var(--xr-border-color);\n",
       "}\n",
       "\n",
       ".xr-group-box-contents {\n",
       "  grid-column-start: 3;\n",
       "  padding-bottom: 4px;\n",
       "}\n",
       "\n",
       ".xr-group-box-contents > label::before {\n",
       "  content: \"📂\";\n",
       "  padding-right: 0.3em;\n",
       "}\n",
       "\n",
       ".xr-group-box-contents > input:checked + label::before {\n",
       "  content: \"📁\";\n",
       "}\n",
       "\n",
       ".xr-group-box-contents > input:checked + label {\n",
       "  padding-bottom: 0px;\n",
       "}\n",
       "\n",
       ".xr-group-box-contents > input:checked ~ .xr-sections {\n",
       "  display: none;\n",
       "}\n",
       "\n",
       ".xr-group-box-contents > input + label > span {\n",
       "  display: none;\n",
       "}\n",
       "\n",
       ".xr-group-box-ellipsis {\n",
       "  font-size: 1.4em;\n",
       "  font-weight: 900;\n",
       "  color: var(--xr-font-color2);\n",
       "  letter-spacing: 0.15em;\n",
       "  cursor: default;\n",
       "}\n",
       "\n",
       ".xr-array-wrap {\n",
       "  grid-column: 1 / -1;\n",
       "  display: grid;\n",
       "  grid-template-columns: 20px auto;\n",
       "}\n",
       "\n",
       ".xr-array-wrap > label {\n",
       "  grid-column: 1;\n",
       "  vertical-align: top;\n",
       "}\n",
       "\n",
       ".xr-preview {\n",
       "  color: var(--xr-font-color3);\n",
       "}\n",
       "\n",
       ".xr-array-preview,\n",
       ".xr-array-data {\n",
       "  padding: 0 5px !important;\n",
       "  grid-column: 2;\n",
       "}\n",
       "\n",
       ".xr-array-data,\n",
       ".xr-array-in:checked ~ .xr-array-preview {\n",
       "  display: none;\n",
       "}\n",
       "\n",
       ".xr-array-in:checked ~ .xr-array-data,\n",
       ".xr-array-preview {\n",
       "  display: inline-block;\n",
       "}\n",
       "\n",
       ".xr-dim-list {\n",
       "  display: inline-block !important;\n",
       "  list-style: none;\n",
       "  padding: 0 !important;\n",
       "  margin: 0;\n",
       "}\n",
       "\n",
       ".xr-dim-list li {\n",
       "  display: inline-block;\n",
       "  padding: 0;\n",
       "  margin: 0;\n",
       "}\n",
       "\n",
       ".xr-dim-list:before {\n",
       "  content: \"(\";\n",
       "}\n",
       "\n",
       ".xr-dim-list:after {\n",
       "  content: \")\";\n",
       "}\n",
       "\n",
       ".xr-dim-list li:not(:last-child):after {\n",
       "  content: \",\";\n",
       "  padding-right: 5px;\n",
       "}\n",
       "\n",
       ".xr-has-index {\n",
       "  font-weight: bold;\n",
       "}\n",
       "\n",
       ".xr-var-list,\n",
       ".xr-var-item {\n",
       "  display: contents;\n",
       "}\n",
       "\n",
       ".xr-var-item > div,\n",
       ".xr-var-item label,\n",
       ".xr-var-item > .xr-var-name span {\n",
       "  background-color: var(--xr-background-color-row-even);\n",
       "  border-color: var(--xr-background-color-row-odd);\n",
       "  margin-bottom: 0;\n",
       "  padding-top: 2px;\n",
       "}\n",
       "\n",
       ".xr-var-item > .xr-var-name:hover span {\n",
       "  padding-right: 5px;\n",
       "}\n",
       "\n",
       ".xr-var-list > li:nth-child(odd) > div,\n",
       ".xr-var-list > li:nth-child(odd) > label,\n",
       ".xr-var-list > li:nth-child(odd) > .xr-var-name span {\n",
       "  background-color: var(--xr-background-color-row-odd);\n",
       "  border-color: var(--xr-background-color-row-even);\n",
       "}\n",
       "\n",
       ".xr-var-name {\n",
       "  grid-column: 1;\n",
       "}\n",
       "\n",
       ".xr-var-dims {\n",
       "  grid-column: 2;\n",
       "}\n",
       "\n",
       ".xr-var-dtype {\n",
       "  grid-column: 3;\n",
       "  text-align: right;\n",
       "  color: var(--xr-font-color2);\n",
       "}\n",
       "\n",
       ".xr-var-preview {\n",
       "  grid-column: 4;\n",
       "}\n",
       "\n",
       ".xr-index-preview {\n",
       "  grid-column: 2 / 5;\n",
       "  color: var(--xr-font-color2);\n",
       "}\n",
       "\n",
       ".xr-var-name,\n",
       ".xr-var-dims,\n",
       ".xr-var-dtype,\n",
       ".xr-preview,\n",
       ".xr-attrs dt {\n",
       "  white-space: nowrap;\n",
       "  overflow: hidden;\n",
       "  text-overflow: ellipsis;\n",
       "  padding-right: 10px;\n",
       "}\n",
       "\n",
       ".xr-var-name:hover,\n",
       ".xr-var-dims:hover,\n",
       ".xr-var-dtype:hover,\n",
       ".xr-attrs dt:hover {\n",
       "  overflow: visible;\n",
       "  width: auto;\n",
       "  z-index: 1;\n",
       "}\n",
       "\n",
       ".xr-var-attrs,\n",
       ".xr-var-data,\n",
       ".xr-index-data {\n",
       "  display: none;\n",
       "  border-top: 2px dotted var(--xr-background-color);\n",
       "  padding-bottom: 20px !important;\n",
       "  padding-top: 10px !important;\n",
       "}\n",
       "\n",
       ".xr-var-attrs-in + label,\n",
       ".xr-var-data-in + label,\n",
       ".xr-index-data-in + label {\n",
       "  padding: 0 1px;\n",
       "}\n",
       "\n",
       ".xr-var-attrs-in:checked ~ .xr-var-attrs,\n",
       ".xr-var-data-in:checked ~ .xr-var-data,\n",
       ".xr-index-data-in:checked ~ .xr-index-data {\n",
       "  display: block;\n",
       "}\n",
       "\n",
       ".xr-var-data > table {\n",
       "  float: right;\n",
       "}\n",
       "\n",
       ".xr-var-data > pre,\n",
       ".xr-index-data > pre,\n",
       ".xr-var-data > table > tbody > tr {\n",
       "  background-color: transparent !important;\n",
       "}\n",
       "\n",
       ".xr-var-name span,\n",
       ".xr-var-data,\n",
       ".xr-index-name div,\n",
       ".xr-index-data,\n",
       ".xr-attrs {\n",
       "  padding-left: 25px !important;\n",
       "}\n",
       "\n",
       ".xr-attrs,\n",
       ".xr-var-attrs,\n",
       ".xr-var-data,\n",
       ".xr-index-data {\n",
       "  grid-column: 1 / -1;\n",
       "}\n",
       "\n",
       "dl.xr-attrs {\n",
       "  padding: 0;\n",
       "  margin: 0;\n",
       "  display: grid;\n",
       "  grid-template-columns: 125px auto;\n",
       "}\n",
       "\n",
       ".xr-attrs dt,\n",
       ".xr-attrs dd {\n",
       "  padding: 0;\n",
       "  margin: 0;\n",
       "  float: left;\n",
       "  padding-right: 10px;\n",
       "  width: auto;\n",
       "}\n",
       "\n",
       ".xr-attrs dt {\n",
       "  font-weight: normal;\n",
       "  grid-column: 1;\n",
       "}\n",
       "\n",
       ".xr-attrs dt:hover span {\n",
       "  display: inline-block;\n",
       "  background: var(--xr-background-color);\n",
       "  padding-right: 10px;\n",
       "}\n",
       "\n",
       ".xr-attrs dd {\n",
       "  grid-column: 2;\n",
       "  white-space: pre-wrap;\n",
       "  word-break: break-all;\n",
       "}\n",
       "\n",
       ".xr-icon-database,\n",
       ".xr-icon-file-text2,\n",
       ".xr-no-icon {\n",
       "  display: inline-block;\n",
       "  vertical-align: middle;\n",
       "  width: 1em;\n",
       "  height: 1.5em !important;\n",
       "  stroke-width: 0;\n",
       "  stroke: currentColor;\n",
       "  fill: currentColor;\n",
       "}\n",
       "\n",
       ".xr-var-attrs-in:checked + label > .xr-icon-file-text2,\n",
       ".xr-var-data-in:checked + label > .xr-icon-database,\n",
       ".xr-index-data-in:checked + label > .xr-icon-database {\n",
       "  color: var(--xr-font-color0);\n",
       "  filter: drop-shadow(1px 1px 5px var(--xr-font-color2));\n",
       "  stroke-width: 0.8px;\n",
       "}\n",
       "</style><pre class='xr-text-repr-fallback'>&lt;xarray.DataArray &#x27;normal-758a57bb75ed3127c94212c3e83a5ed8&#x27; (batch: 1000)&gt; Size: 8kB\n",
       "77.07 77.61 76.86 77.53 77.99 77.22 76.9 ... 77.92 78.03 77.96 76.6 77.31 76.26\n",
       "Dimensions without coordinates: batch</pre><div class='xr-wrap' style='display:none'><div class='xr-header'><div class='xr-obj-type'>xarray.DataArray</div><div class='xr-obj-name'>&#x27;normal-758a57bb75ed3127c94212c3e83a5ed8&#x27;</div><ul class='xr-dim-list'><li><span>batch</span>: 1000</li></ul></div><ul class='xr-sections'><li class='xr-section-item'><div class='xr-array-wrap'><input id='section-74798c8e-b829-4c3c-99f8-3f476cc5db19' class='xr-array-in' type='checkbox' ><label for='section-74798c8e-b829-4c3c-99f8-3f476cc5db19' title='Show/hide data repr'><svg class='icon xr-icon-database'><use xlink:href='#icon-database'></use></svg></label><div class='xr-array-preview xr-preview'><span>77.07 77.61 76.86 77.53 77.99 77.22 ... 78.03 77.96 76.6 77.31 76.26</span></div><div class='xr-array-data'><pre>array([77.07151123, 77.61059979, 76.86360161, 77.52666541, 77.98533059,\n",
       "       77.22236039, 76.89694279, 76.5113661 , 77.04069162, 77.0440349 ,\n",
       "       76.38611509, 77.30566883, 77.47694384, 76.69118445, 76.3648431 ,\n",
       "       78.65136421, 75.48230116, 77.29719159, 76.69989052, 77.53014796,\n",
       "       78.12950831, 76.58109187, 76.1735932 , 77.57211317, 76.31369742,\n",
       "       76.28008721, 76.00930925, 77.69946061, 78.47248892, 76.30594417,\n",
       "       76.74855535, 77.23211165, 77.09690595, 77.58629277, 77.25236896,\n",
       "       77.29720169, 77.70320198, 77.29337301, 77.51158852, 77.00113117,\n",
       "       77.61088289, 76.56355944, 76.66462471, 78.35684907, 77.62398703,\n",
       "       77.15455699, 77.42167832, 76.20287995, 75.56872634, 77.17161862,\n",
       "       77.6550905 , 77.3271054 , 77.00974287, 78.52692208, 76.62211391,\n",
       "       76.63936546, 77.72529412, 75.9304276 , 77.94115594, 76.28197731,\n",
       "       77.86032362, 77.34556766, 78.28347133, 78.25610074, 76.4990315 ,\n",
       "       77.05555913, 77.7762863 , 76.83277802, 75.61555854, 77.93083141,\n",
       "       76.35444611, 76.40108947, 77.10925417, 76.21775021, 76.50717149,\n",
       "       78.62424972, 78.41599829, 77.53805218, 76.29311862, 77.23899266,\n",
       "       77.17823025, 76.23308925, 77.10158854, 78.06501202, 77.81267966,\n",
       "       75.9464039 , 77.37662775, 76.22437435, 76.93255673, 77.05499553,\n",
       "       78.42548546, 76.63240287, 78.43203628, 77.27066793, 76.95095667,\n",
       "       77.6427735 , 78.23776876, 77.98181894, 76.55946822, 76.59540871,\n",
       "...\n",
       "       77.55074864, 77.08245593, 76.05782067, 77.52499048, 78.07990445,\n",
       "       76.95812276, 76.09941659, 77.63809857, 75.79134748, 77.27308766,\n",
       "       77.26043898, 76.87723759, 76.20732289, 76.91405135, 78.67692647,\n",
       "       76.95501719, 77.22709319, 76.90604196, 77.1620702 , 75.95007881,\n",
       "       76.99028891, 77.16984169, 77.32234885, 76.23063511, 76.81834779,\n",
       "       76.15704279, 75.61472124, 78.17831097, 76.7443434 , 76.50639742,\n",
       "       77.62027759, 75.975392  , 76.57688827, 77.47909537, 77.36364272,\n",
       "       77.42555982, 76.23375548, 76.81310853, 77.87934912, 76.94914194,\n",
       "       77.22891481, 77.04548581, 77.41784027, 76.99567393, 75.91612618,\n",
       "       77.92538295, 77.87590131, 75.91280074, 77.01375182, 76.84293853,\n",
       "       77.29023139, 75.74463574, 76.52104818, 76.9560999 , 76.93174665,\n",
       "       75.94148142, 75.8301277 , 75.99251482, 77.34913468, 78.26571634,\n",
       "       77.86725889, 76.94211159, 76.92407092, 76.70451861, 77.12471401,\n",
       "       77.26577934, 76.06551508, 76.75742254, 76.21261178, 76.68219503,\n",
       "       76.59003924, 76.59220068, 75.72158433, 76.32204662, 77.53317453,\n",
       "       77.35568039, 78.39380204, 78.01003637, 76.57906144, 76.44788913,\n",
       "       76.16136169, 76.78692734, 77.15260584, 76.97373823, 78.16022121,\n",
       "       77.12782383, 76.94807175, 76.84872826, 76.92376664, 75.95237816,\n",
       "       76.72480888, 76.89486946, 77.19022507, 77.66279427, 77.91773289,\n",
       "       78.02838772, 77.96117143, 76.60281147, 77.3129846 , 76.25648097])</pre></div></div></li></ul></div></div>"
      ],
      "text/plain": [
       "<xarray.DataArray 'normal-758a57bb75ed3127c94212c3e83a5ed8' (batch: 1000)> Size: 8kB\n",
       "77.07 77.61 76.86 77.53 77.99 77.22 76.9 ... 77.92 78.03 77.96 76.6 77.31 76.26\n",
       "Dimensions without coordinates: batch"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "linalg.norm(matrix_batch, dims=[\"dim1\", \"dim2\"], dask=\"parallelized\").compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d03e25c-c592-414e-869f-efd2bf033dc1",
   "metadata": {},
   "source": [
    ":::{seealso}\n",
    "{ref}`Xarray documentation <xarray:dask>`\n",
    "\n",
    "[Dask examples](https://examples.dask.org/xarray.html)\n",
    ":::"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "c336f026-9460-4ee8-bbd4-d3343ae4ddcb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Last updated: Mon, 20 Jul 2026\n",
      "\n",
      "Python implementation: CPython\n",
      "Python version       : 3.14.6\n",
      "IPython version      : 9.15.0\n",
      "\n",
      "dask           : 2026.7.1\n",
      "xarray         : 2026.7.0\n",
      "xarray_einstats: 0.11.0\n",
      "\n",
      "Watermark: 2.6.0\n",
      "\n"
     ]
    }
   ],
   "source": [
    "%load_ext watermark\n",
    "%watermark -n -u -v -iv -w"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb4ffe5e-7f62-4523-8974-d5db0f26aaa0",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.14.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}