Add xbeam.normalize_chunks() and update xbeam.Dataset docstrings

PiperOrigin-RevId: 813800457

Author: shoyer

docs/api.md

@@ -103,4 +103,5 @@ guarantees.
     Dataset.mean
     Dataset.head
     Dataset.pipe
+    normalize_chunks
 ```

docs/high-level.ipynb

@@ -24,28 +24,6 @@
         "```"
       ]
     },
-    {
-      "metadata": {
-        "id": "Ht5100QIRpyd",
-        "tags": [
-          "hide-cell"
-        ]
-      },
-      "cell_type": "code",
-      "source": [
-        "# small formatting improvements\n",
-        "import contextlib\n",
-        "\n",
-        "@contextlib.contextmanager\n",
-        "def print_error():\n",
-        "  try:\n",
-        "    yield\n",
-        "  except Exception as e:\n",
-        "    print(f'{type(e).__name__}: {e}')"
-      ],
-      "outputs": [],
-      "execution_count": 2
-    },
     {
       "metadata": {
         "id": "EkH9Na4ezuTi"
@@ -254,11 +232,13 @@
       },
       "cell_type": "code",
       "source": [
-        "with print_error():\n",
+        "try:\n",
         "  (\n",
         "      xbeam.Dataset.from_zarr('example_data.zarr')\n",
         "      .map_blocks(lambda ds: ds.compute())  # load into memory\n",
-        "  )"
+        "  )\n",
+        "except Exception as e:\n",
+        "  print(f'{type(e).__name__}: {e}')"
       ],
       "outputs": [],
       "execution_count": 9
@@ -290,7 +270,7 @@
       },
       "cell_type": "markdown",
       "source": [
-        "In other situations, you might want to perform an operation that returns something other than an `xarray.Dataset`, e.g., to write all chunks as individual files to disk. In these situations, you can switch to the lower-level Xarray-Beam [data model](data-model), and use raw Beam operations:"
+        "Sometimes, your computation doesn't fit into the ``map_blocks`` paradigm because you don't want to create `xarray.Dataset` objects. For these cases, you can switch to the lower-level Xarray-Beam [data model](data-model), and use raw Beam operations:"
       ]
     },
     {

xarray_beam/__init__.py

@@ -13,43 +13,46 @@
 # limitations under the License.
 """Public API for Xarray-Beam."""
 
-# pylint: disable=g-multiple-import
+# Note: import <name> as <name> is required for names to be exported.
+# See PEP 484 & https://github.com/jax-ml/jax/issues/7570
+# pylint: disable=g-multiple-import,useless-import-alias,g-importing-member
 from xarray_beam._src.combiners import (
-    Mean,
-    MeanCombineFn,
+    Mean as Mean,
+    MeanCombineFn as MeanCombineFn,
 )
 from xarray_beam._src.core import (
-    Key,
-    DatasetToChunks,
-    ValidateEachChunk,
-    offsets_to_slices,
-    validate_chunk
+    Key as Key,
+    DatasetToChunks as DatasetToChunks,
+    ValidateEachChunk as ValidateEachChunk,
+    offsets_to_slices as offsets_to_slices,
+    validate_chunk as validate_chunk,
 )
 from xarray_beam._src.dataset import (
-    Dataset,
+    Dataset as Dataset,
+    normalize_chunks as normalize_chunks,
 )
 from xarray_beam._src.rechunk import (
-    ConsolidateChunks,
-    ConsolidateVariables,
-    SplitChunks,
-    SplitVariables,
-    Rechunk,
-    consolidate_chunks,
-    consolidate_variables,
-    consolidate_fully,
-    split_chunks,
-    split_variables,
-    in_memory_rechunk,
+    ConsolidateChunks as ConsolidateChunks,
+    ConsolidateVariables as ConsolidateVariables,
+    SplitChunks as SplitChunks,
+    SplitVariables as SplitVariables,
+    Rechunk as Rechunk,
+    consolidate_chunks as consolidate_chunks,
+    consolidate_variables as consolidate_variables,
+    consolidate_fully as consolidate_fully,
+    split_chunks as split_chunks,
+    split_variables as split_variables,
+    in_memory_rechunk as in_memory_rechunk,
 )
 from xarray_beam._src.zarr import (
-    open_zarr,
-    make_template,
-    replace_template_dims,
-    setup_zarr,
-    validate_zarr_chunk,
-    write_chunk_to_zarr,
-    ChunksToZarr,
-    DatasetToZarr,
+    open_zarr as open_zarr,
+    make_template as make_template,
+    replace_template_dims as replace_template_dims,
+    setup_zarr as setup_zarr,
+    validate_zarr_chunk as validate_zarr_chunk,
+    write_chunk_to_zarr as write_chunk_to_zarr,
+    ChunksToZarr as ChunksToZarr,
+    DatasetToZarr as DatasetToZarr,
 )
 
-__version__ = '0.10.1'  # automatically synchronized to pyproject.toml
+__version__ = '0.10.2'  # automatically synchronized to pyproject.toml