From 0f7771baed15d2c079e42822c80ba1732bc879be Mon Sep 17 00:00:00 2001 From: thanos Date: Thu, 8 Oct 2026 19:22:43 -0400 Subject: [PATCH] Everything passes: - Tests: 1,679 tests, 70 properties and 155 doctests, plus the performance tests. - Checks: format, compile with warnings as errors, credo, doctor, sobelow, docs, and Dialyzer. I rebuilt your stale local Dialyzer cache, which is a gitignored build file. - Livebooks: all 17 shipped livebooks run against this branch's code. MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit To show the new tests check something real, I ran them against the released 1.3.0 code. They fail exactly where behaviour changed: the dtype error, crop speed, the caller's random state, and the livebook harness. Code - Edge-chunk crop (F2): it now copies one contiguous run per row. A 999×1000 edge chunk takes under 1 ms instead of 305 ms. The results match to_tensor/2 on 1D through 4D arrays with edges on every axis. While making this change I found that the review's "duplicate shape computation" point was wrong: the stream's metadata.shape is the nominal chunk shape, not the cropped one. Using it broke every edge chunk, so the code still computes the cropped shape from the chunk bounds. - Shuffled batching (F3): - Batches are grouped in a single pass, which removes the quadratic slowdown. - The caller's random state is untouched, with or without a seed. The default seed was quietly using the caller's random generator, which I only caught by testing. - A given seed produces exactly the same order as 1.3.0, pinned by tests. - A short final batch that drop_remainder will drop is no longer read first. - Smaller fixes: an unsupported dtype now gives one error instead of one per chunk (F13), the unreachable clause is gone (F12), and one shared helper handles :names and :backend. That last change means to_tensor_chunked/3 now honours :names too. Livebooks and tooling - Codecs livebook (F5): it crashed on gzip, which is a Zarr v3 codec only, and then on compressor: nil, which needs to be :none. Both are fixed, and it runs end to end. - Axon notebook (F6): the second feature was an exact copy of the first plus 0.05, so the "right weights" the notebook described eren't learnable. It now uses an independent second feature, zero starting weights, Adam and Axon 0.8's model-state API. It prints weights: [2.0, -0.0], bias: [1.0], predictions: [1.0, 3.0], the same on every run, with no deprecation warning. - mix livebook.test (F7): it now runs cells one at a time like Livebook, names the failing cell, and keeps the end of long output where the error is. A new --local option tests against the checkout, and a new CI job runs every shipped livebook that way. Docs - Explorer scope (F1): the CHANGELOG, release notes and ROADMAP now say Explorer streaming was planned and dropped, with links to ex_arrow and ex_graphblas. The plan file's statuses are updated, with Explorer marked cancelled. - Speed claims: the "5-10× faster, 400-800 MB/s" claims in the Nx guide and module docs didn't hold up. Measured on your M1 Max with an 8 MB array, to_tensor/2 takes 75 ms uncompressed and 117 ms with the default zlib. I replaced the claims with those labelled numbers. - Remaining review itms: the reduction example uses include_missing and on_error: :halt (F9), the stale shuffle example is fixed (F8), and the pin (F10), comma (F11), headings (F15), broken bold (F16), new long dashes and CHANGELOG wording (F14) are corrected. Beyond the review, found while fixing: - Six guides said :gzip "always works". It only works for Zarr v3 arrays, so they now say so. - compressor: nil also appeared in two livebook formance.md; those are fixed too. --- .github/workflows/ci.yml | 37 +++++ CHANGELOG.md | 32 +++- README.md | 4 +- docs/ROADMAP.md | 14 +- docs/guides/core_concepts.md | 2 +- docs/guides/getting_started.md | 4 +- docs/guides/glossary.md | 2 +- docs/guides/installation.md | 4 +- docs/guides/nx_integration.md | 39 +++-- docs/guides/overview.md | 4 +- docs/guides/performance.md | 2 +- docs/guides/troubleshooting.md | 2 +- .../01_01_first_zarr_array.livemd | 2 +- .../01_04_codecs_and_pipelines.livemd | 17 +- .../03_nx_ml/03_03_training_from_zarr.livemd | 83 +++++----- docs/livebooks/nx_streaming.livemd | 6 +- docs/livebooks/scripts/00_01_welcome.exs | 2 +- .../scripts/01_02_metadata_and_chunks.exs | 2 +- .../scripts/01_04_codecs_and_pipelines.exs | 4 +- docs/release_notes_v1_3_0.md | 12 +- lib/ex_zarr/nx.ex | 149 ++++++++---------- lib/ex_zarr/nx/data_loader.ex | 146 ++++++----------- lib/mix/tasks/livebook.test.ex | 90 ++++++++--- test/ex_zarr/nx/data_loader_test.exs | 57 +++++++ test/ex_zarr/nx_test.exs | 142 +++++++++++++++++ test/mix/tasks/livebook_test_task_test.exs | 51 ++++++ 26 files changed, 618 insertions(+), 291 deletions(-) create mode 100644 test/mix/tasks/livebook_test_task_test.exs diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 3b81ea0..36fd221 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -182,6 +182,43 @@ jobs: echo "Checking for undocumented public functions..." mix docs --warnings-as-errors + livebooks: + name: Livebooks + runs-on: ubuntu-latest + env: + MIX_ENV: dev + + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set up Elixir + uses: erlef/setup-beam@v1 + with: + elixir-version: "1.19.5" + otp-version: "28.1" + + - name: Restore dependencies cache + uses: actions/cache@v4 + with: + path: | + deps + _build + key: ${{ runner.os }}-mix-livebooks-${{ hashFiles('**/mix.lock') }} + restore-keys: | + ${{ runner.os }}-mix-livebooks- + + - name: Install dependencies + run: mix deps.get + + # Runs every shipped notebook cell by cell against this checkout, so a + # change that breaks a livebook fails here rather than after release. + - name: Run livebooks + run: >- + mix livebook.test --local --max 2 --timeout 900 + --paths "docs/livebooks/0[1-6]*/*.livemd" + --paths docs/livebooks/nx_streaming.livemd + python-interop: name: Python Interop (zarr ${{ matrix.zarr }}) runs-on: ubuntu-latest diff --git a/CHANGELOG.md b/CHANGELOG.md index e063a75..332a1a3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,9 +7,39 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Fixed +- `ExZarr.Nx.stream_chunk_tensors/2` crops padded edge chunks with one + contiguous copy per row instead of one copy per element. A 999x1000 float32 + edge chunk went from about 300 ms to under 1 ms. +- `stream_chunk_tensors/2` on an array whose dtype has no Nx type now yields a + single `{:error, reason}` and reads nothing, instead of one error per chunk. +- `DataLoader.shuffled_batch_stream/3` and `paired_shuffled_batch_stream/4` + group shuffled indices in one pass. Slicing the index list per batch was + quadratic in the sample count (about 1 s of indexing at 200,000 samples). + A short final batch is no longer read when `drop_remainder: true` drops it. +- Shuffling no longer reseeds the caller's `:rand` state, with or without + `:seed`. A given `:seed` still produces the same order as 1.3.0. +- `to_tensor_chunked/3` now applies `:names`, like `to_tensor/2`. +- Livebooks: `01_04_codecs_and_pipelines` no longer uses `:gzip` with Zarr v2 + (gzip is a Zarr v3 codec; v2 uses `:zlib`) or `compressor: nil` (use + `:none`). `03_03_training_from_zarr` uses an independent second feature, + zero-initialized weights and Adam, so it reliably learns the line, and uses + Axon 0.8's `Axon.ModelState`. +- Guides no longer say `:gzip` always works; it is accepted for Zarr v3 arrays. + +### Changed +- `mix livebook.test` evaluates notebook cells one at a time with bindings and + environment carried forward, like Livebook, names the failing cell, and keeps + the end of long output. `--local` runs notebooks against this checkout + instead of the Hex release. CI runs `mix livebook.test --local`. + ## [1.3.0] - 2026-10-08 -### Data Science Interop +Data science interop for dense arrays. + +Not included, although planned for 1.3: Explorer direct streaming. Columnar +data and Explorer integration belong in [ex_arrow](https://hex.pm/packages/ex_arrow); +Binsparse belongs in [ex_graphblas](https://hex.pm/packages/ex_graphblas). ### Added - `ExZarr.Nx.stream_chunk_tensors/2` maps `Array.stream_chunks/2` to Nx tensors. diff --git a/README.md b/README.md index 62c7a47..df92477 100644 --- a/README.md +++ b/README.md @@ -833,8 +833,8 @@ Practical examples demonstrating real-world usage: - Intuitive slicing by name instead of index - Real-world examples (climate, medical imaging) - Validation and best practices -- **[Using Nx](docs/guides/getting_started.md#using-nx)** — what Nx is, when it fits, and how tensors and Zarr chunks exchange bytes -- **[Nx Integration](docs/guides/nx_integration.md)** — conversion API, chunk streams, and DataLoader +- **[Using Nx](docs/guides/getting_started.md#using-nx)**: what Nx is, when it fits, and how tensors and Zarr chunks exchange bytes +- **[Nx Integration](docs/guides/nx_integration.md)**: conversion API, chunk streams, and DataLoader - Converting between Nx tensors and Zarr arrays - Machine learning workflows - Streaming large arrays diff --git a/docs/ROADMAP.md b/docs/ROADMAP.md index 673bc69..b486a80 100644 --- a/docs/ROADMAP.md +++ b/docs/ROADMAP.md @@ -1,6 +1,6 @@ # Roadmap -## v1.2.0 — Zarr 3.1 Interoperability & Range-Aware Cloud I/O +## v1.2.0 - Zarr 3.1 Interoperability & Range-Aware Cloud I/O Released 2026-10-04. Backward compatible with v1.1 (rewrite pre-1.2 private shards; OTP 26+). @@ -16,7 +16,7 @@ Released 2026-10-04. Backward compatible with v1.1 (rewrite pre-1.2 private shar --- -## v1.1.0 — BEAM-Native Streaming +## v1.1.0 - BEAM-Native Streaming Released 2026-06-12. Backward compatible with v1.0. @@ -26,18 +26,18 @@ Released 2026-06-12. Backward compatible with v1.0. --- -## v1.3.0 — Data Science Interop +## v1.3.0 - Data Science Interop Released 2026-10-08. - [x] `ExZarr.Nx.stream_chunk_tensors/2` (stored chunks → tensors, edge crops) - [x] DataLoader documented against its real shuffle (full index list) - [x] Livebook MVP: metadata, Zarr↔Nx, minibatches, Axon toy train -- [x] Arrow/Explorer stay in ex_arrow; Binsparse stays in ex_graphblas +- [ ] ~~Explorer direct streaming~~ Dropped from 1.3: columnar data and Explorer belong in [ex_arrow](https://hex.pm/packages/ex_arrow), and Binsparse in [ex_graphblas](https://hex.pm/packages/ex_graphblas) --- -## v1.4.0 (Planned) — Performance & Packaging +## v1.4.0 (Planned) - Performance & Packaging - [ ] Async codec pipeline - [ ] Vendored/static codec libraries @@ -46,14 +46,14 @@ Released 2026-10-08. --- -## v2.0.0 (Future) — Distributed Processing +## v2.0.0 (Future) - Distributed Processing - [ ] Multi-node chunk processing - [ ] Cross-node telemetry --- -## Deferred from v1.2 (explicit) +## Deferred (explicit) - Unified retry layer across cloud SDKs - Native multi-range / coalescing heuristics diff --git a/docs/guides/core_concepts.md b/docs/guides/core_concepts.md index 2eeb754..c2b3d9b 100644 --- a/docs/guides/core_concepts.md +++ b/docs/guides/core_concepts.md @@ -357,7 +357,7 @@ Codecs transform chunk data before storage. This includes compression (reduce si **Compression codecs** (reduce size): - `:zlib` - Standard compression (Erlang built-in, always available) -- `:gzip` - Compatible with gzip tools (Erlang built-in) +- `:gzip` - Compatible with gzip tools (Erlang built-in; Zarr v3 arrays only, use `:zlib` for v2) - `:zstd` - Fast compression with excellent ratio - `:lz4` - Very fast compression, lower ratio - `:blosc` - Shuffle filters plus an inner compressor; strong on numeric data diff --git a/docs/guides/getting_started.md b/docs/guides/getting_started.md index abc5842..0bdb4c0 100644 --- a/docs/guides/getting_started.md +++ b/docs/guides/getting_started.md @@ -471,7 +471,7 @@ Skip it when the access pattern is a few scalar lookups. A Zarr read loads a chu Skip it for `{:bf, 16}`, `{:f, 16}`, and complex tensors. Those types are not in the Zarr dtypes ExZarr stores. Keep float32 on disk and cast after `to_tensor/2`. -### Quick Guide to Nx functionality +### Nx functions at a glance | Function | What you get | | --- | --- | @@ -495,7 +495,7 @@ array The livebooks under `docs/livebooks/03_nx_ml/` and `docs/livebooks/nx_streaming.livemd` run these calls. The rest of the API is in the [Nx integration guide](nx_integration.md). -### Deep Dive in how data is stored in Nx versus Zarr and how we exchange data between the two +### How data moves between Nx and Zarr An Nx tensor is one buffer. For `Nx.iota({6, 4}, type: {:f, 32})` that buffer is `6 * 4 * 4` bytes, row-major (C order), with the shape stored beside the binary. Names and the backend are metadata in the tensor struct. Nothing is compressed. Nothing is split into files. diff --git a/docs/guides/glossary.md b/docs/guides/glossary.md index e29110b..f270c19 100644 --- a/docs/guides/glossary.md +++ b/docs/guides/glossary.md @@ -349,7 +349,7 @@ Note: `<` indicates little-endian byte order (standard on most modern systems). | Codec | Python (numcodecs) | ExZarr/Elixir | Availability | |-------|-------------------|---------------|--------------| | Zlib | `numcodecs.Zlib` | `:zlib` | Always (Erlang built-in) | -| Gzip | `numcodecs.GZip` | `:gzip` | Always (Erlang built-in) | +| Gzip | `numcodecs.GZip` | `:gzip` | Zarr v3 only (Erlang built-in) | | Zstd | `numcodecs.Zstd` | `:zstd` | Via ExCodecs | | LZ4 | `numcodecs.LZ4` | `:lz4` | Via ExCodecs | | Blosc | `numcodecs.Blosc` | `:blosc` | Via ExCodecs (Blosc1 chunks) | diff --git a/docs/guides/installation.md b/docs/guides/installation.md index 7e435aa..ddd9042 100644 --- a/docs/guides/installation.md +++ b/docs/guides/installation.md @@ -209,8 +209,8 @@ failed (offline build, proxy) or the platform has no precompiled binary. # config/config.exs config :rustler_precompiled, :force_build, ex_codecs: true ``` -3. Meanwhile `:zlib` and `:gzip` keep working; they use Erlang's built-in - `:zlib`. +3. Meanwhile `:zlib` (and `:gzip` on Zarr v3 arrays) keep working; they use + Erlang's built-in `:zlib`. ### Issue: Permission Denied Errors diff --git a/docs/guides/nx_integration.md b/docs/guides/nx_integration.md index 5e8c7a2..3df613a 100644 --- a/docs/guides/nx_integration.md +++ b/docs/guides/nx_integration.md @@ -4,7 +4,7 @@ This guide shows how to integrate ExZarr with Nx (Numerical Elixir) for numerica For what Nx is, when to use the bridge, and how tensor bytes relate to Zarr chunks, start with [Using Nx](getting_started.md#using-nx) in the getting started guide. This page is the API: conversion, chunk streams, and DataLoader. -** Performance Note:** ExZarr provides an optimized `ExZarr.Nx` module with **5-10x faster** conversion compared to manual approaches. Always use `ExZarr.Nx` for best performance. +**Performance note:** `ExZarr.Nx` converts through one binary instead of nested lists, so it skips the cost of building and walking millions of list cells. Measured numbers are in [Optimized Conversion](#optimized-conversion-recommended). ## Table of Contents @@ -42,7 +42,7 @@ tensor = Nx.iota({1000, 1000}) ) ``` -**Performance:** 10-20ms for 8MB (400-800 MB/s) +**Performance:** an 8 MB float64 array converts in tens of milliseconds without compression; codec choice dominates. See the table under [Optimized Conversion](#optimized-conversion-recommended). ## Nx and ExZarr Overview @@ -89,12 +89,19 @@ ExZarr provides the `ExZarr.Nx` module for efficient conversion using direct bin ### Performance Comparison -| Approach | Time (8MB) | Throughput | Status | -|----------|-----------|------------|---------| -| **`ExZarr.Nx` (optimized)** | **10-20ms** | **400-800 MB/s** | **Recommended** | -| Nested tuples (legacy) | 80-150ms | 50-100 MB/s | Deprecated | +Measured on an Apple M1 Max: a `{1000, 1000}` float64 array (8 MB), chunks +`{100, 100}`, memory storage, median of seven runs. -**Speedup: 5-10x faster** +| Compressor | `to_tensor/2` | `from_tensor/2` | +|------------|---------------|-----------------| +| `:none` | 75 ms | 50 ms | +| `:zstd` | 87 ms | 98 ms | +| `:zlib` (the default) | 117 ms | 640 ms | + +The legacy nested-list path does the same storage work and adds the list +conversion on top: about 400 ms for `Nx.to_list/1` on the write side and about +100 ms for `Nx.tensor/2` on the read side for this array. Compression is the +largest cost, so choose the codec before tuning anything else. ### Converting ExZarr to Nx @@ -397,9 +404,12 @@ trained_state = ### Chunk-aligned reduction ```elixir +# include_missing: never-written chunks count as the fill value. +# on_error: :halt: a chunk that fails to read stops the sum instead of +# silently leaving it out. sums = array - |> ExZarr.Nx.stream_chunk_tensors(concurrency: 4, on_error: :skip) + |> ExZarr.Nx.stream_chunk_tensors(concurrency: 4, include_missing: true, on_error: :halt) |> Enum.map(fn {:ok, tensor} -> Nx.sum(tensor) |> Nx.to_number() end) total = Enum.sum(sums) @@ -723,7 +733,7 @@ result = tensor **Problem**: Conversion requires copying data (binary → tensor). -**Impact**: ~10-20ms overhead per 8MB. +**Impact**: one copy of the data; for 8 MB, tens of milliseconds (see the table under [Optimized Conversion](#optimized-conversion-recommended)). **Workarounds:** 1. Process in chunks (amortize overhead) @@ -805,7 +815,7 @@ end ### WARNING: Deprecated: Nested Tuple Conversion -**Note:** This approach is **5-10x slower** than `ExZarr.Nx`. Only use for compatibility with old code. +**Note:** This approach is slower than `ExZarr.Nx` because it builds nested lists (see the measurements above). Only use it for compatibility with old code.
Click to expand legacy nested tuple approach @@ -817,7 +827,7 @@ defmodule ExZarr.Nx.LegacyHelpers do @moduledoc """ Legacy conversion helpers using nested tuples. - WARNING: DEPRECATED: Use ExZarr.Nx module instead (5-10x faster). + WARNING: DEPRECATED: Use the ExZarr.Nx module instead. """ def nested_list_to_tuple(list) when is_list(list) do @@ -839,7 +849,7 @@ end #### Legacy Write Pattern ```elixir -# WARNING: SLOW: 80-150ms for 8MB +# WARNING: SLOW: adds a nested-list conversion on top of the storage work tensor = Nx.iota({1000, 1000}) {:ok, array} = ExZarr.create( @@ -864,7 +874,7 @@ data = tensor #### Legacy Read Pattern ```elixir -# WARNING: SLOW: 80-150ms for 8MB +# WARNING: SLOW: adds a nested-list conversion on top of the storage work {:ok, array} = ExZarr.open(path: "/data/array") {:ok, data} = ExZarr.Array.get_slice(array, @@ -885,13 +895,12 @@ tensor = data ExZarr provides first-class Nx integration via the `ExZarr.Nx` module: -**5-10x faster** than manual conversion (400-800 MB/s vs 50-100 MB/s) +**No nested lists** - conversion goes through one binary **Simple API** - Single function calls for conversion **Full type support** - All 10 standard numeric types **Chunked processing** - Constant memory for large arrays **Backend agnostic** - Works with CPU, EXLA, Torchx **ML-ready** - Efficient batch loading, checkpointing -**Production-tested** - Used in real-world workflows **Next steps:** - Try the example: `elixir examples/nx_optimized_conversion.exs` diff --git a/docs/guides/overview.md b/docs/guides/overview.md index 04c9f1f..dd1c83e 100644 --- a/docs/guides/overview.md +++ b/docs/guides/overview.md @@ -111,8 +111,8 @@ ExCodecs publishes binaries for macOS (Apple Silicon and Intel), Linux (x86_64 and aarch64, glibc and musl) and Windows (x86_64). If its NIF cannot load, ExZarr does not silently switch codecs: `ExZarr.Codecs.available_codecs/0` lists only `:none` and `:zlib`, and compressing with another codec returns -`{:error, {:unsupported_codec, codec}}` or a compression error. `:zlib` and -`:gzip` always work. +`{:error, {:unsupported_codec, codec}}` or a compression error. `:zlib` always +works, and so does `:gzip` on Zarr v3 arrays. ### NIF Safety diff --git a/docs/guides/performance.md b/docs/guides/performance.md index 145db8a..de18699 100644 --- a/docs/guides/performance.md +++ b/docs/guides/performance.md @@ -315,7 +315,7 @@ Data integrity critical (checksums required)? → Add crc32c codec to pipeline Data already compressed (images, video)? - → Use no compression (compressor: nil) + → Use no compression (compressor: :none) ``` ### Compression Level Tuning diff --git a/docs/guides/troubleshooting.md b/docs/guides/troubleshooting.md index 4d3a376..568eccd 100644 --- a/docs/guides/troubleshooting.md +++ b/docs/guides/troubleshooting.md @@ -1011,7 +1011,7 @@ python3 -c "import zarr; z = zarr.open('/path/to/array', mode='r'); print(z.info ```elixir # Python zarr always supports zlib/gzip compressor: :zlib - # Or + # Or, for a Zarr v3 array compressor: :gzip ``` diff --git a/docs/livebooks/01_core_zarr/01_01_first_zarr_array.livemd b/docs/livebooks/01_core_zarr/01_01_first_zarr_array.livemd index 7639562..6a8420b 100644 --- a/docs/livebooks/01_core_zarr/01_01_first_zarr_array.livemd +++ b/docs/livebooks/01_core_zarr/01_01_first_zarr_array.livemd @@ -2,7 +2,7 @@ ```elixir Mix.install([ - {:ex_zarr, "~> 1.2"}, + {:ex_zarr, "~> 1.3"}, # {:ex_zarr, path: Path.join(__DIR__, "../../..")} ]) ``` diff --git a/docs/livebooks/01_core_zarr/01_04_codecs_and_pipelines.livemd b/docs/livebooks/01_core_zarr/01_04_codecs_and_pipelines.livemd index bb9f64a..ca84f8a 100644 --- a/docs/livebooks/01_core_zarr/01_04_codecs_and_pipelines.livemd +++ b/docs/livebooks/01_core_zarr/01_04_codecs_and_pipelines.livemd @@ -16,7 +16,7 @@ Compression is critical for Zarr arrays — it reduces storage costs, speeds up * Zarr v2 compressor model (single-stage) * Zarr v3 codec pipelines (multi-stage) -* Common codecs: gzip, zstd, zlib, blosc +* Common codecs: zstd, zlib, blosc (v2 and v3), gzip (v3 only) * Compression ratio vs speed trade-offs * Choosing codecs for different data types * Measuring compression effectiveness @@ -48,8 +48,8 @@ graph LR **Simple and effective:** For most use cases, a single compressor is enough. v2 supports: -* `gzip` - Widely compatible, moderate speed -* `zlib` - Similar to gzip, slightly different format +* `zlib` - Widely compatible, moderate speed (the v2 counterpart of v3 `gzip`) +* `gzip` - Zarr v3 only; ExZarr rejects it as a v2 compressor * `zstd` - Faster, better ratios, modern choice * `blosc` - Optimized for numerical data, multiple algorithms * `lz4` - Extremely fast, lower compression ratio @@ -87,7 +87,7 @@ These are called "array-to-array" or "array-to-bytes" codecs, followed by a "byt Let's create the same logical array with different v2 compressors and compare. ```elixir -compressors = [:zstd, :gzip, :zlib, :blosc] +compressors = [:zstd, :zlib, :blosc] arrays = Enum.map(compressors, fn comp -> @@ -162,7 +162,7 @@ end) **Expected patterns:** * `zstd`: Fast, good compression -* `gzip/zlib`: Slower, decent compression +* `zlib`: Slower, decent compression * `blosc`: Very fast, excellent for numerical data * `lz4`: Fastest, lower compression ratio @@ -282,8 +282,7 @@ graph TD * **zstd (default):** Best general-purpose choice. Fast, good compression, widely supported. * **blosc:** Excellent for numerical arrays (satellite data, simulation results). Use with zstd or lz4 backend. -* **gzip:** Maximum compatibility. Slower, but readable by any tool. -* **zlib:** Similar to gzip, slight differences in format. +* **zlib (v2) / gzip (v3):** Maximum compatibility. Slower, but readable by any tool. ExZarr accepts `gzip` only for Zarr v3 arrays. * **lz4:** When speed is critical and compression ratio is secondary (e.g., temporary intermediate data). **For Zarr v3:** Start with zstd. Add byte shuffle if working with numerical data. Consider transpose for multidimensional arrays with non-uniform access patterns. @@ -390,7 +389,7 @@ Compression isn't free. Let's measure the overhead. shape: {500, 500}, chunks: {100, 100}, dtype: :float32, - compressor: nil, + compressor: :none, storage: :memory ) @@ -473,7 +472,7 @@ Smaller datasets are easier to share, download, and distribute. Compressed array 4. **blosc excels for numerical data** — built-in shuffling improves compression 5. **Compression ratios depend on data** — structured data compresses better 6. **Codec overhead is usually worth it** — I/O savings outweigh CPU cost -7. **Choose based on priority** — speed (lz4), compression (zstd high level), compatibility (gzip) +7. **Choose based on priority** — speed (lz4), compression (zstd high level), compatibility (zlib for v2, gzip for v3) ## What's Next diff --git a/docs/livebooks/03_nx_ml/03_03_training_from_zarr.livemd b/docs/livebooks/03_nx_ml/03_03_training_from_zarr.livemd index c33cb89..0aadd5e 100644 --- a/docs/livebooks/03_nx_ml/03_03_training_from_zarr.livemd +++ b/docs/livebooks/03_nx_ml/03_03_training_from_zarr.livemd @@ -5,7 +5,8 @@ Mix.install([ {:ex_zarr, "~> 1.3"}, # {:ex_zarr, path: Path.join(__DIR__, "../../..")}, {:nx, "~> 0.7"}, - {:axon, "~> 0.7"} # Axon is installed only in this notebook. ExZarr does not depend on it. + # Axon is installed only in this notebook. ExZarr does not depend on it. + {:axon, "~> 0.8"} ]) ``` @@ -13,30 +14,34 @@ Mix.install([ This notebook trains a one-layer model on data that lives in ExZarr arrays. -The labels are a straight line of the first feature: for sample `i`, `y = 2 * x0 + 1`. The second feature is present in the array and is not used by that formula. The model has to learn which column matters. +The labels are a straight line of the first feature: for sample `i`, `y = 2 * x0 + 1`. The second feature is stored in the array but plays no part in that formula, and it does not follow the first feature, so the model has to learn to ignore it. -Training never loads all 64 rows at once. `ExZarr.Nx.DataLoader` reads 16 samples at a time, which is also the chunk size, and Axon takes one gradient step per batch. +The training loop never holds more than one batch: `ExZarr.Nx.DataLoader` reads 16 samples at a time, which is also the chunk size, and Axon takes one gradient step per batch. ## Write the line into two arrays There are 64 samples. Each feature row is two float32 numbers: * column 0: `i * 0.05` -* column 1: `(i + 1) * 0.05` +* column 1: `rem(i * 7, 13) * 0.1`, a repeating pattern unrelated to column 0 and to the target -The target for that row is `2 * (i * 0.05) + 1`, which is `2 * x0 + 1`. Sample 0 is features `[0.0, 0.05]` and target `1.0`. Sample 1 is `[0.05, 0.10]` and target `1.1`. +The target for that row is `2 * (i * 0.05) + 1`, which is `2 * x0 + 1`. Sample 0 is features `[0.0, 0.0]` and target `1.0`. Sample 1 is `[0.05, 0.7]` and target `1.1`. -Both arrays use chunks of 16 along the sample axis, so the file is four chunks. The batch size below is also 16, so each training step reads exactly one chunk. +Both arrays use chunks of 16 along the sample axis, so each array is four chunks. The batch size below is also 16, so each training step reads exactly one chunk. ```elixir alias ExZarr.Array n = 64 -features = for(i <- 0..(n - 1), j <- 0..1, into: <<>>, do: <<(i + j) * 0.05::float-little-32>>) + +features = + for i <- 0..(n - 1), into: <<>> do + <> + end targets = for i <- 0..(n - 1), into: <<>> do - <<(2.0 * (i * 0.05) + 1.0)::float-little-32>> + <<2.0 * (i * 0.05) + 1.0::float-little-32>> end {:ok, x_array} = @@ -48,45 +53,41 @@ targets = :ok = Array.set_slice(x_array, features, start: {0, 0}, stop: {n, 2}) :ok = Array.set_slice(y_array, targets, start: {0, 0}, stop: {n, 1}) -{:ok, x0} = ExZarr.Nx.to_tensor(x_array) -{:ok, y0} = ExZarr.Nx.to_tensor(y_array) +{:ok, first_two} = Array.get_slice(x_array, start: {0, 0}, stop: {2, 2}) %{ - x_shape: Nx.shape(x0), - y_shape: Nx.shape(y0), - first_row: Nx.to_list(x0[0]), - first_target: y0[0][0] |> Nx.to_number() + x_shape: x_array.shape, + y_shape: y_array.shape, + first_rows: first_two |> Nx.from_binary({:f, 32}) |> Nx.reshape({2, 2}) |> Nx.to_list() } ``` -The map should show `x_shape: {64, 2}`, `y_shape: {64, 1}`, `first_row: [0.0, 0.05]`, and `first_target: 1.0`. +The map shows `x_shape: {64, 2}`, `y_shape: {64, 1}`, and the first two feature rows, `[[0.0, 0.0], [0.05, 0.7]]` (float32 values print with extra digits, such as `0.05000000074505806` and `0.699999988079071`). ## One dense layer, batches of 16 `Axon.input("features", shape: {nil, 2})` is the placeholder for a batch. `nil` means the number of rows can change. `2` is the two columns in `x_array`. -`Axon.dense(1)` is one output number per row: `y_hat = w0 * x0 + w1 * x1 + b`. That is the whole model. Training changes `w0`, `w1`, and `b`. A good fit has `w0` near 2, `w1` near 0, and `b` near 1, because the labels ignore column 1. +`Axon.dense(1)` is one output number per row: `y_hat = w0 * x0 + w1 * x1 + b`. That is the whole model. Training changes `w0`, `w1`, and `b`. The right answer is `w0 = 2`, `w1 = 0`, and `b = 1`, because the labels ignore column 1. -`paired_batch_stream/4` reads the same row range from both arrays, so feature row `i` stays with target `i`. `drop_remainder: true` would drop a short final batch. Here 64 ÷ 16 is exact, so every batch is `{16, 2}` features and `{16, 1}` targets. +The weights start at zero (`kernel_initializer: :zeros`, `bias_initializer: :zeros`) instead of random values. With the batches always in the same order, every run of this notebook prints the same numbers. -Axon wants each item as `{inputs, targets}`. The inputs are a map whose key matches the input name `"features"`. `Nx.as_type/2` makes the dtype explicit before the training loop compiles. +`paired_batch_stream/4` reads the same row range from both arrays, so feature row `i` stays with target `i`. `drop_remainder: true` would drop a short final batch. Here 64 / 16 is exact, so every batch is `{16, 2}` features and `{16, 1}` targets. + +Axon wants each item as `{inputs, targets}`. The inputs are a map whose key matches the input name `"features"`. ```elixir model = Axon.input("features", shape: {nil, 2}) - |> Axon.dense(1) + |> Axon.dense(1, kernel_initializer: :zeros, bias_initializer: :zeros) batches = ExZarr.Nx.DataLoader.paired_batch_stream(x_array, y_array, 16, drop_remainder: true) - |> Stream.map(fn {:ok, {x, y}} -> - {Nx.as_type(x, {:f, 32}), Nx.as_type(y, {:f, 32})} - end) - |> Stream.map(fn {x, y} -> {%{"features" => x}, y} end) + |> Stream.map(fn {:ok, {x, y}} -> {%{"features" => x}, y} end) {first_inputs, first_targets} = batches |> Enum.take(1) |> hd() %{ - model: model, batch_x: Nx.shape(first_inputs["features"]), batch_y: Nx.shape(first_targets) } @@ -94,33 +95,35 @@ batches = You should see batch shapes `{16, 2}` and `{16, 1}`. -## Twelve gradient steps, then two predictions - -`Axon.Loop.trainer/3` builds a loop that minimizes mean squared error with stochastic gradient descent. Mean squared error is the average of `(y_hat - y)²` over the batch. Each step nudges the three parameters to make that number smaller. +## Train, then check the weights and two predictions -`epochs: 3` walks the data three times. `iterations: 4` stops each pass after four batches, which is the whole dataset (64 / 16). That is 12 updates, starting from Axon's random initial weights. This is a demonstration that the batches are readable, not a claim that the line is fully learned. +`Axon.Loop.trainer/3` builds a loop that minimizes mean squared error, the average of `(y_hat - y)²` over the batch. The Adam optimizer with a learning rate of 0.1 adjusts the three parameters after every batch to make that number smaller. -`Axon.Loop.run/4` returns the trained parameters. Some Axon versions wrap them in `%Axon.Loop.State{}`. The `case` unwraps that so `Axon.predict/3` receives the model state either way. +`epochs: 60` walks the four batches sixty times: 240 updates, each reading one chunk from each array. `Axon.Loop.run/4` starts from an empty `Axon.ModelState` and returns the trained one. The two probe rows are chosen from the formula, not from the stored array: * `[0.0, 0.0]` has `x0 = 0`, so the line says `y = 1` * `[1.0, 0.0]` has `x0 = 1`, so the line says `y = 3` -After a dozen steps the printed numbers should sit closer to `1` and `3` than a random guess, and they will not match those targets exactly. - ```elixir -trainer = Axon.Loop.trainer(model, :mean_squared_error, :sgd) +trainer = + Axon.Loop.trainer(model, :mean_squared_error, Polaris.Optimizers.adam(learning_rate: 0.1), + log: 0 + ) -trained = Axon.Loop.run(trainer, batches, %{}, epochs: 3, iterations: 4) +model_state = Axon.Loop.run(trainer, batches, Axon.ModelState.empty(), epochs: 60) -model_state = - case trained do - %Axon.Loop.State{step_state: %{model_state: state}} -> state - state -> state - end +round_list = fn tensor -> tensor |> Nx.to_flat_list() |> Enum.map(&Float.round(&1, 2)) end + +predictions = + Axon.predict(model, model_state, %{"features" => Nx.tensor([[0.0, 0.0], [1.0, 0.0]])}) -Axon.predict(model, model_state, %{"features" => Nx.tensor([[0.0, 0.0], [1.0, 0.0]])}) +%{ + weights: round_list.(model_state.data["dense_0"]["kernel"]), + bias: round_list.(model_state.data["dense_0"]["bias"]), + predictions: round_list.(predictions) +} ``` -The result is a `{2, 1}` tensor. Row 0 is the prediction for `x0 = 0`. Row 1 is the prediction for `x0 = 1`. +The result is `weights: [2.0, 0.0]` (the second weight may print as `-0.0`), `bias: [1.0]`, and `predictions: [1.0, 3.0]`: the model learned the line and learned to ignore column 1, reading the data one chunk at a time from the Zarr arrays. diff --git a/docs/livebooks/nx_streaming.livemd b/docs/livebooks/nx_streaming.livemd index 1b74d42..bf9a48c 100644 --- a/docs/livebooks/nx_streaming.livemd +++ b/docs/livebooks/nx_streaming.livemd @@ -4,13 +4,13 @@ Two ways to turn a Zarr array into Nx tensors: -- `ExZarr.Nx.stream_chunk_tensors/2` — one tensor per stored chunk (`Array.stream_chunks/2`) -- `ExZarr.Nx.DataLoader` — fixed-size batches along the first axis (`Array.get_slice/2`) +- `ExZarr.Nx.stream_chunk_tensors/2`: one tensor per stored chunk (`Array.stream_chunks/2`) +- `ExZarr.Nx.DataLoader`: fixed-size batches along the first axis (`Array.get_slice/2`) ```elixir Mix.install([ {:ex_zarr, "~> 1.3"}, - # {:ex_zarr, path: Path.join(__DIR__, "../..")} + # {:ex_zarr, path: Path.join(__DIR__, "../..")}, {:nx, "~> 0.7"}, {:kino, "~> 0.13"} ]) diff --git a/docs/livebooks/scripts/00_01_welcome.exs b/docs/livebooks/scripts/00_01_welcome.exs index a5cd2e8..8ae5f60 100644 --- a/docs/livebooks/scripts/00_01_welcome.exs +++ b/docs/livebooks/scripts/00_01_welcome.exs @@ -546,7 +546,7 @@ IO.puts("Full array shape: #{inspect(Nx.shape(full))}") shape: {100, 100}, chunks: {50, 50}, dtype: :float64, - compressor: nil, + compressor: :none, storage: :memory ) diff --git a/docs/livebooks/scripts/01_02_metadata_and_chunks.exs b/docs/livebooks/scripts/01_02_metadata_and_chunks.exs index 6f72cb6..03ff6be 100644 --- a/docs/livebooks/scripts/01_02_metadata_and_chunks.exs +++ b/docs/livebooks/scripts/01_02_metadata_and_chunks.exs @@ -358,7 +358,7 @@ IO.inspect(metadata_v3, limit: :infinity) # * `shape`: Tuple of dimensions # * `chunks`: Tuple of chunk dimensions # * `dtype`: Atom like `:float32`, `:int64` -# * `compressor`: Atom like `:zstd`, `:gzip` +# * `compressor`: Atom like `:zstd`, `:zlib` (`:gzip` on v3 only) # * `fill_value`: Default value # * `zarr_format`: Integer (2 or 3) diff --git a/docs/livebooks/scripts/01_04_codecs_and_pipelines.exs b/docs/livebooks/scripts/01_04_codecs_and_pipelines.exs index fc63bcb..8b01931 100644 --- a/docs/livebooks/scripts/01_04_codecs_and_pipelines.exs +++ b/docs/livebooks/scripts/01_04_codecs_and_pipelines.exs @@ -83,7 +83,7 @@ alias ExZarr.Gallery.{Pack, SampleData, Metrics} # Let's create the same logical array with different v2 compressors and compare. -compressors = [:zstd, :gzip, :zlib, :blosc] +compressors = [:zstd, :zlib, :blosc] arrays = Enum.map(compressors, fn comp -> @@ -374,7 +374,7 @@ Result: 10-30x compression for typical climate data shape: {500, 500}, chunks: {100, 100}, dtype: :float32, - compressor: nil, + compressor: :none, storage: :memory ) diff --git a/docs/release_notes_v1_3_0.md b/docs/release_notes_v1_3_0.md index dc9a237..4fb63d1 100644 --- a/docs/release_notes_v1_3_0.md +++ b/docs/release_notes_v1_3_0.md @@ -7,7 +7,13 @@ small livebook set you can run from the repo. ExZarr remains a dense chunked array and group store. Columnar batches belong in [ex_arrow](https://hex.pm/packages/ex_arrow). Binsparse groups belong in -ex_graphblas. +[ex_graphblas](https://hex.pm/packages/ex_graphblas). + +### Scope change + +Explorer direct streaming was planned for 1.3 and is not included. Explorer +works on columnar data, which is ex_arrow's job, so it will not be added to +ExZarr. ## Highlights @@ -24,8 +30,8 @@ array |> Enum.map(fn {:ok, tensor} -> Nx.sum(tensor) end) ``` -`:concurrency`, `:ordered`, `:on_error`, `:filter`, and `:include_missing` -are the streaming options. `:backend` and `:names` apply to the tensors. +`:concurrency`, `:ordered`, `:timeout`, `:on_error`, `:filter`, and +`:include_missing` are the streaming options. `:backend` and `:names` apply to the tensors. ### Batches diff --git a/lib/ex_zarr/nx.ex b/lib/ex_zarr/nx.ex index 9b8421d..38df394 100644 --- a/lib/ex_zarr/nx.ex +++ b/lib/ex_zarr/nx.ex @@ -7,10 +7,11 @@ defmodule ExZarr.Nx do ## Performance - Direct binary conversion is **5-10x faster** than nested tuple approach: - - - **Optimized (this module)**: 10-20ms for 8MB (400-800 MB/s) - - **Nested tuples**: 80-150ms for 8MB (50-100 MB/s) + Conversion goes through one binary instead of nested lists. Chunk + decompression is usually the largest cost: for an 8 MB float64 array in + memory storage on an Apple M1 Max, `to_tensor/2` took about 75 ms + uncompressed and 117 ms with the default zlib. Building nested lists instead + adds roughly 100 to 400 ms for the same array. ## Usage @@ -112,10 +113,10 @@ defmodule ExZarr.Nx do ## Performance - Conversion time scales linearly with array size: - - 1 MB: ~2ms - - 10 MB: ~15ms - - 100 MB: ~150ms + Conversion time scales linearly with array size and is dominated by chunk + decompression. Measured on an Apple M1 Max for an 8 MB float64 array in + memory storage: about 75 ms uncompressed, 87 ms with zstd, and 117 ms with + the default zlib. For arrays larger than available RAM, use `to_tensor_chunked/3` instead. @@ -136,25 +137,11 @@ defmodule ExZarr.Nx do def to_tensor(%Array{} = array, opts \\ []) do with {:ok, binary} <- Array.to_binary(array), {:ok, nx_type} <- zarr_to_nx_type(array.metadata.dtype) do - # Create tensor from binary (1D) - tensor = Nx.from_binary(binary, nx_type) - - # Reshape to original array shape - tensor = Nx.reshape(tensor, array.metadata.shape) - - # Apply axis names if provided - tensor = - case Keyword.get(opts, :names) do - nil -> tensor - names -> Nx.rename(tensor, names) - end - - # Transfer to backend if specified tensor = - case Keyword.get(opts, :backend) do - nil -> tensor - backend -> Nx.backend_transfer(tensor, backend) - end + binary + |> Nx.from_binary(nx_type) + |> Nx.reshape(array.metadata.shape) + |> apply_tensor_opts(opts) {:ok, tensor} end @@ -263,8 +250,17 @@ defmodule ExZarr.Nx do `:metadata` is forced on internally so edge-chunk shapes are known; the stream still yields `{:ok, tensor}` or `{:error, reason}`, not chunk maps. - `:on_error` behaves like `stream_chunks/2`: `:skip` drops failed chunks, - `:halt` raises `ExZarr.StreamError`. + `:on_error` behaves like `stream_chunks/2`: `:skip` (the default) drops failed + chunks without yielding anything for them, `:halt` raises `ExZarr.StreamError`, + and a two-arity function is called with the failing index. Use `:halt` when a + silently missing chunk would make a result wrong, for example a total. + + Without `:include_missing`, only chunks that exist in storage are yielded. + Pass `include_missing: true` to also get never-written chunks, filled with the + array's fill value. + + When the array dtype has no Nx equivalent, the stream yields a single + `{:error, reason}` and reads nothing. ## Examples @@ -278,11 +274,16 @@ defmodule ExZarr.Nx do Enumerable.t({:ok, Nx.Tensor.t()} | {:error, term()}) def stream_chunk_tensors(%Array{} = array, opts \\ []) do {tensor_opts, stream_opts} = Keyword.split(opts, [:backend, :names]) - stream_opts = Keyword.put(stream_opts, :metadata, true) - array - |> Array.stream_chunks(stream_opts) - |> Stream.map(&chunk_event_to_tensor(array, &1, tensor_opts)) + case zarr_to_nx_type(array.metadata.dtype) do + {:ok, nx_type} -> + array + |> Array.stream_chunks(Keyword.put(stream_opts, :metadata, true)) + |> Stream.map(&chunk_event_to_tensor(array, nx_type, &1, tensor_opts)) + + {:error, _} = error -> + [error] + end end @doc """ @@ -326,18 +327,11 @@ defmodule ExZarr.Nx do Stream.map(chunk_ranges, fn {start, stop} -> with {:ok, binary} <- Array.get_slice(array, start: start, stop: stop), {:ok, nx_type} <- zarr_to_nx_type(array.metadata.dtype) do - # Calculate chunk shape - chunk_shape = calculate_chunk_shape(start, stop) - - # Create tensor - tensor = Nx.from_binary(binary, nx_type) |> Nx.reshape(chunk_shape) - - # Apply backend transfer if specified tensor = - case Keyword.get(opts, :backend) do - nil -> tensor - backend -> Nx.backend_transfer(tensor, backend) - end + binary + |> Nx.from_binary(nx_type) + |> Nx.reshape(calculate_chunk_shape(start, stop)) + |> apply_tensor_opts(opts) {:ok, tensor} end @@ -513,26 +507,15 @@ defmodule ExZarr.Nx do |> List.to_tuple() end - defp chunk_event_to_tensor(array, %{index: index, data: data}, opts) + # The event's metadata.shape is the nominal chunk shape; the in-bounds shape + # comes from the chunk bounds. + defp chunk_event_to_tensor(array, nx_type, %{index: index, data: data}, opts) when is_binary(data) do - binary_to_chunk_tensor(array, index, data, opts) - end - - defp chunk_event_to_tensor(array, {index, data}, opts) when is_binary(data) do - binary_to_chunk_tensor(array, index, data, opts) - end - - defp chunk_event_to_tensor(_array, other, _opts) do - {:error, {:unexpected_chunk_event, other}} - end - - defp binary_to_chunk_tensor(array, index, data, opts) do {start_coords, stop_coords} = Array.get_chunk_bounds(array, index) shape = calculate_chunk_shape(start_coords, stop_coords) - nominal = array.chunks - with {:ok, nx_type} <- zarr_to_nx_type(array.metadata.dtype), - {:ok, payload} <- crop_chunk_binary(data, shape, nominal, Array.itemsize(array)) do + with {:ok, payload} <- + crop_chunk_binary(data, shape, array.chunks, Array.itemsize(array)) do tensor = payload |> Nx.from_binary(nx_type) @@ -561,35 +544,39 @@ defmodule ExZarr.Nx do end end + # Copies the in-bounds region of a padded chunk. In C order, the dimensions + # after the last cropped one span the whole chunk, so from that dimension on + # each region row is one contiguous run. Cropping only the first dimension is + # a single binary_part/3. defp extract_logical_chunk(data, shape, nominal, itemsize) do - strides = chunk_strides(nominal) + dims = Tuple.to_list(shape) + nominal_dims = Tuple.to_list(nominal) - shape - |> tuple_coordinates() + {_, split} = + Enum.zip(dims, nominal_dims) + |> Enum.with_index() + |> Enum.filter(fn {{dim, nom}, _} -> dim != nom end) + |> List.last() + + {lead_dims, [run_first | _]} = Enum.split(dims, split) + lead_strides = nominal |> chunk_strides() |> Enum.take(split) + run_bytes = Enum.reduce(Enum.drop(nominal_dims, split + 1), run_first, &*/2) * itemsize + + lead_dims + |> Enum.map(&Range.new(0, &1 - 1, 1)) + |> cartesian_product() |> Enum.map(fn coord -> - offset = flat_offset(coord, strides) * itemsize - binary_part(data, offset, itemsize) + binary_part(data, flat_offset(coord, lead_strides) * itemsize, run_bytes) end) |> IO.iodata_to_binary() end defp chunk_strides(shape) do - dims = Tuple.to_list(shape) - - dims - |> Enum.reverse() - |> Enum.reduce({[], 1}, fn dim, {acc, stride} -> - {[stride | acc], stride * dim} - end) - |> elem(0) - end - - defp tuple_coordinates(shape) do shape |> Tuple.to_list() - |> Enum.map(&(&1 - 1)) - |> Enum.map(&Range.new(0, &1)) - |> cartesian_product() + |> Enum.reverse() + |> Enum.reduce({[], 1}, fn dim, {acc, stride} -> {[stride | acc], stride * dim} end) + |> elem(0) end defp flat_offset(coord, strides) do @@ -601,7 +588,11 @@ defmodule ExZarr.Nx do shape |> Tuple.to_list() |> Enum.reduce(1, &*/2) end - defp apply_tensor_opts(tensor, opts) do + @doc false + # Applies the `:names` and `:backend` options shared by every conversion in + # this module and in `ExZarr.Nx.DataLoader`. + @spec apply_tensor_opts(Nx.Tensor.t(), keyword()) :: Nx.Tensor.t() + def apply_tensor_opts(tensor, opts) do tensor = case Keyword.get(opts, :names) do nil -> tensor diff --git a/lib/ex_zarr/nx/data_loader.ex b/lib/ex_zarr/nx/data_loader.ex index bc6fe95..26ce6dc 100644 --- a/lib/ex_zarr/nx/data_loader.ex +++ b/lib/ex_zarr/nx/data_loader.ex @@ -175,9 +175,13 @@ defmodule ExZarr.Nx.DataLoader do ## Shuffling Strategy Builds the full sample-index list, shuffles it, then loads each batch with - `get_slice/2` (one read per sample). Memory for the index list is + `get_slice/2` (one read per sample). Memory and time for the index list are O(number of samples). `:shuffle_buffer_size` is ignored. + Shuffling seeds the process's `:rand` generator with `:seed` and restores the + caller's previous `:rand` state afterwards, so it does not change other random + numbers in the calling process. + ## Options - `:seed` - Random seed for reproducibility (default: random) @@ -210,10 +214,9 @@ defmodule ExZarr.Nx.DataLoader do end) end - # Large shuffle buffer for better randomization - array - |> ExZarr.Nx.DataLoader.shuffled_batch_stream(32, shuffle_buffer_size: 1000) - |> Enum.to_list() + # The same seed gives the same order on every run + first = ExZarr.Nx.DataLoader.shuffled_batch_stream(array, 32, seed: 42) |> Enum.to_list() + ^first = ExZarr.Nx.DataLoader.shuffled_batch_stream(array, 32, seed: 42) |> Enum.to_list() """ @spec shuffled_batch_stream(ExZarr.Array.t(), pos_integer(), [batch_option()]) :: @@ -222,42 +225,12 @@ defmodule ExZarr.Nx.DataLoader do when is_integer(batch_size) and batch_size > 0 do shape = array.metadata.shape num_samples = elem(shape, 0) - seed = Keyword.get(opts, :seed, :rand.uniform(1_000_000)) - drop_remainder = Keyword.get(opts, :drop_remainder, false) - - # Shuffle indices - indices = shuffle_indices(num_samples, seed) + seed = Keyword.get_lazy(opts, :seed, &random_seed/0) - # Calculate number of batches - num_batches = - if drop_remainder do - div(num_samples, batch_size) - else - div(num_samples + batch_size - 1, batch_size) - end - - Stream.unfold(0, fn batch_idx -> - if batch_idx < num_batches do - start_idx = batch_idx * batch_size - end_idx = min(start_idx + batch_size, num_samples) - actual_batch_size = end_idx - start_idx - - # Get shuffled indices for this batch - batch_indices = Enum.slice(indices, start_idx, actual_batch_size) - - # Load samples at shuffled indices - result = load_samples_at_indices(array, batch_indices, opts) - - # Skip if dropping remainder and batch is incomplete - if drop_remainder and actual_batch_size < batch_size do - nil - else - {result, batch_idx + 1} - end - else - nil - end - end) + num_samples + |> shuffle_indices(seed) + |> index_batches(batch_size, Keyword.get(opts, :drop_remainder, false)) + |> Stream.map(&load_samples_at_indices(array, &1, opts)) end @doc """ @@ -391,42 +364,15 @@ defmodule ExZarr.Nx.DataLoader do "Got #{num_features} features and #{num_labels} labels." end - seed = Keyword.get(opts, :seed, :rand.uniform(1_000_000)) - drop_remainder = Keyword.get(opts, :drop_remainder, false) - - # Shuffle indices (same for both arrays) - indices = shuffle_indices(num_features, seed) + seed = Keyword.get_lazy(opts, :seed, &random_seed/0) - num_batches = - if drop_remainder do - div(num_features, batch_size) - else - div(num_features + batch_size - 1, batch_size) - end - - Stream.unfold(0, fn batch_idx -> - if batch_idx < num_batches do - start_idx = batch_idx * batch_size - end_idx = min(start_idx + batch_size, num_features) - actual_batch_size = end_idx - start_idx - - # Skip if dropping remainder and batch is incomplete - if drop_remainder and actual_batch_size < batch_size do - nil - else - # Get shuffled indices for this batch - batch_indices = Enum.slice(indices, start_idx, actual_batch_size) - - # Load both batches with same indices - with {:ok, x_batch} <- load_samples_at_indices(features, batch_indices, opts), - {:ok, y_batch} <- load_samples_at_indices(labels, batch_indices, opts) do - {{:ok, {x_batch, y_batch}}, batch_idx + 1} - else - {:error, reason} -> {{:error, reason}, batch_idx + 1} - end - end - else - nil + num_features + |> shuffle_indices(seed) + |> index_batches(batch_size, Keyword.get(opts, :drop_remainder, false)) + |> Stream.map(fn batch_indices -> + with {:ok, x_batch} <- load_samples_at_indices(features, batch_indices, opts), + {:ok, y_batch} <- load_samples_at_indices(labels, batch_indices, opts) do + {:ok, {x_batch, y_batch}} end end) end @@ -483,7 +429,7 @@ defmodule ExZarr.Nx.DataLoader do tensor = Nx.from_binary(binary, nx_type) |> Nx.reshape(batch_shape) # Apply options - tensor = apply_tensor_options(tensor, opts) + tensor = ExZarrNx.apply_tensor_opts(tensor, opts) {:ok, tensor} end @@ -523,7 +469,7 @@ defmodule ExZarr.Nx.DataLoader do stacked = Nx.stack(tensors) # Apply options - stacked = apply_tensor_options(stacked, opts) + stacked = ExZarrNx.apply_tensor_opts(stacked, opts) {:ok, stacked} else @@ -532,27 +478,39 @@ defmodule ExZarr.Nx.DataLoader do end end - defp shuffle_indices(num_samples, seed) do - # Create range and shuffle - _ = :rand.seed(:exsss, {seed, seed * 2, seed * 3}) - 0..(num_samples - 1) |> Enum.to_list() |> Enum.shuffle() + # Default seed when :seed is not given. Drawn from :crypto so the caller's + # :rand state is not advanced. + defp random_seed do + <> = :crypto.strong_rand_bytes(4) + rem(value, 1_000_000) + 1 end - defp apply_tensor_options(tensor, opts) do - # Apply axis names if provided - tensor = - case Keyword.get(opts, :names) do - nil -> tensor - names -> Nx.rename(tensor, names) + # Same algorithm and seeding as 1.3.0, so a seed gives the same order, but the + # caller's :rand state is restored afterwards instead of being overwritten. + defp shuffle_indices(num_samples, seed) do + previous = :rand.export_seed() + + try do + _ = :rand.seed(:exsss, {seed, seed * 2, seed * 3}) + 0..(num_samples - 1)//1 |> Enum.to_list() |> Enum.shuffle() + after + case previous do + :undefined -> Process.delete(:rand_seed) + state -> :rand.seed(state) end + end + end - # Transfer to backend if specified - tensor = - case Keyword.get(opts, :backend) do - nil -> tensor - backend -> Nx.backend_transfer(tensor, backend) - end + # Groups indices into batches in one pass (Enum.slice/3 per batch was + # quadratic in the number of samples). A short final batch is dropped + # before it is read when drop_remainder is set. + defp index_batches(indices, batch_size, drop_remainder) do + batches = Stream.chunk_every(indices, batch_size) - tensor + if drop_remainder do + Stream.filter(batches, &(length(&1) == batch_size)) + else + batches + end end end diff --git a/lib/mix/tasks/livebook.test.ex b/lib/mix/tasks/livebook.test.ex index 45528e1..2af10e9 100644 --- a/lib/mix/tasks/livebook.test.ex +++ b/lib/mix/tasks/livebook.test.ex @@ -6,6 +6,15 @@ defmodule Mix.Tasks.Livebook.Test do @moduledoc """ Discovers .livemd files, extracts elixir fenced blocks, writes a temp .exs, and runs it. + Cells are evaluated one at a time, in order, carrying variable bindings and + the environment (aliases, imports, requires) forward, the way Livebook does. + Code in later cells is compiled only after earlier cells ran, so a struct from + a package installed by the setup cell's `Mix.install/2` can be used. + + By default the notebooks install `ex_zarr` from Hex, as a reader would. Pass + `--local` to point `{:ex_zarr, "~> ..."}` at this checkout instead, so CI + tests the notebooks against the code being changed. + Why this exists: - No dependency on Livebook - CI-friendly @@ -17,6 +26,7 @@ defmodule Mix.Tasks.Livebook.Test do mix livebook.test --pattern ExZarr --max 4 --timeout 300 mix livebook.test --fail-on-warn mix livebook.test --skip-tag "livebook:test skip" + mix livebook.test --local """ @default_paths ["docs/livebooks/**/*.livemd", "docs/**/*.livemd"] @@ -36,7 +46,8 @@ defmodule Mix.Tasks.Livebook.Test do max: :integer, timeout: :integer, fail_on_warn: :boolean, - skip_tag: :string + skip_tag: :string, + local: :boolean ], aliases: [p: :paths] ) @@ -56,6 +67,7 @@ defmodule Mix.Tasks.Livebook.Test do timeout_sec = Keyword.get(opts, :timeout, @default_timeout_sec) fail_on_warn? = Keyword.get(opts, :fail_on_warn, false) skip_tag = Keyword.get(opts, :skip_tag, "livebook:test skip") + local? = Keyword.get(opts, :local, false) files = patterns @@ -76,7 +88,7 @@ defmodule Mix.Tasks.Livebook.Test do results = files |> Task.async_stream( - fn path -> run_one(path, timeout_sec, fail_on_warn?, skip_tag) end, + fn path -> run_one(path, timeout_sec, fail_on_warn?, skip_tag, local?) end, max_concurrency: max_conc, timeout: (timeout_sec + 30) * 1_000, ordered: true @@ -101,13 +113,13 @@ defmodule Mix.Tasks.Livebook.Test do defp maybe_filter(files, nil), do: files defp maybe_filter(files, pat), do: Enum.filter(files, &String.contains?(&1, pat)) - defp run_one(path, _timeout_sec, fail_on_warn?, skip_tag) do + defp run_one(path, _timeout_sec, fail_on_warn?, skip_tag, local?) do livemd = File.read!(path) if skip_entire_notebook?(livemd, skip_tag) do {:ok, %{path: path, skipped: true}} else - script = __MODULE__.LivemdExtract.to_elixir_script(livemd) + script = __MODULE__.LivemdExtract.to_elixir_script(livemd, local: local?) if blank?(script) do {:error, @@ -151,9 +163,13 @@ defmodule Mix.Tasks.Livebook.Test do Mix.shell().error("❌ FAIL #{path} (#{inspect(status)})\n#{truncate(out)}\n") end + # Keeps the end of the output, where the failing cell and error are. defp truncate(s) when is_binary(s) do max = 12_000 - if byte_size(s) > max, do: binary_part(s, 0, max) <> "\n…(truncated)", else: s + + if byte_size(s) > max, + do: "…(truncated)\n" <> binary_part(s, byte_size(s) - max, max), + else: s end defp tmp_script_path(path) do @@ -189,22 +205,47 @@ defmodule Mix.Tasks.Livebook.Test do # ``` @elixir_fence ~r/```elixir\s*\n(.*?)```/ms - def to_elixir_script(livemd) when is_binary(livemd) do - @elixir_fence - |> Regex.scan(livemd, capture: :all_but_first) - |> List.flatten() - |> Enum.map(&String.trim_trailing/1) - |> Enum.reject(&(String.trim(&1) == "")) - |> Enum.join("\n\n") - |> then(&wrap_header/1) + def to_elixir_script(livemd, opts \\ []) when is_binary(livemd) do + cells = + @elixir_fence + |> Regex.scan(livemd, capture: :all_but_first) + |> List.flatten() + |> Enum.map(&String.trim_trailing/1) + |> Enum.reject(&(String.trim(&1) == "")) + |> Enum.map(&rewrite_paths(&1, Keyword.get(opts, :local, false))) + + if cells == [], do: "", else: runner_script(cells) end - defp wrap_header(body) do - # Replace path-based Mix.install with absolute path to project root + # Each cell is quoted and evaluated with the binding and env left by the + # previous cell, so later cells compile after earlier ones ran. + defp runner_script(cells) do + """ + # Generated from .livemd by mix livebook.test + # NOTE: Only ```elixir fenced blocks are executed, one cell at a time. + cells = #{inspect(cells, limit: :infinity, printable_limit: :infinity)} + + cells + |> Enum.with_index(1) + |> Enum.reduce({[], Code.env_for_eval([])}, fn {code, index}, {binding, env} -> + try do + quoted = Code.string_to_quoted!(code, file: "cell \#{index}") + {_value, binding, env} = Code.eval_quoted_with_env(quoted, binding, env) + {binding, env} + rescue + error -> + IO.puts("livebook cell \#{index} failed:") + reraise error, __STACKTRACE__ + end + end) + """ + end + + defp rewrite_paths(code, local?) do project_root = File.cwd!() - body_fixed = - body + code = + code |> String.replace( ~r/path: Path\.join\(__DIR__, "\.\.\/\.\.\/\.\."\)/, "path: \"#{project_root}\"" @@ -216,12 +257,15 @@ defmodule Mix.Tasks.Livebook.Test do |> String.replace(~r/path: Path\.join\(__DIR__, "\.\."\)/, "path: \"#{project_root}\"") |> String.replace(~r/path: Path\.dirname\(__DIR__\)/, "path: \"#{project_root}\"") - """ - # Generated from .livemd by mix livebook.test - # NOTE: Only ```elixir fenced blocks are executed. - # - #{body_fixed} - """ + if local? do + String.replace( + code, + ~r/^(\s*)\{:ex_zarr, "~> [^"]+"\}/m, + "\\1{:ex_zarr, path: \"#{project_root}\"}" + ) + else + code + end end end end diff --git a/test/ex_zarr/nx/data_loader_test.exs b/test/ex_zarr/nx/data_loader_test.exs index 2a92ddc..e38adeb 100644 --- a/test/ex_zarr/nx/data_loader_test.exs +++ b/test/ex_zarr/nx/data_loader_test.exs @@ -248,6 +248,63 @@ defmodule ExZarr.Nx.DataLoaderTest do end end + describe "shuffled streams: reproducibility and caller state" do + defp int_array(values) do + {:ok, array} = + ExZarr.create(shape: {length(values)}, chunks: {4}, dtype: :int32, storage: :memory) + + data = for v <- values, into: <<>>, do: <> + :ok = ExZarr.Array.set_slice(array, data, start: {0}, stop: {length(values)}) + array + end + + # Orders recorded from v1.3.0, so a seeded epoch is the same across versions. + test "a seed gives the same order as v1.3.0" do + features = int_array(Enum.to_list(0..9)) + labels = int_array(Enum.map(0..9, &(&1 * 10))) + + assert features + |> DataLoader.shuffled_batch_stream(3, seed: 42) + |> Enum.map(fn {:ok, t} -> Nx.to_flat_list(t) end) == + [[1, 3, 9], [4, 5, 7], [2, 6, 8], [0]] + + paired = + features + |> DataLoader.paired_shuffled_batch_stream(labels, 4, seed: 7, drop_remainder: true) + |> Enum.map(fn {:ok, {x, y}} -> {Nx.to_flat_list(x), Nx.to_flat_list(y)} end) + + assert paired == [ + {[1, 4, 3, 9], [10, 40, 30, 90]}, + {[7, 8, 0, 5], [70, 80, 0, 50]} + ] + end + + test "shuffling does not change the caller's :rand state" do + array = int_array(Enum.to_list(0..9)) + + :rand.seed(:exsss, {1, 2, 3}) + before = :rand.export_seed() + _ = array |> DataLoader.shuffled_batch_stream(3, seed: 99) |> Enum.to_list() + _ = array |> DataLoader.shuffled_batch_stream(3) |> Enum.to_list() + assert :rand.export_seed() == before + + Process.delete(:rand_seed) + _ = array |> DataLoader.shuffled_batch_stream(3, seed: 5) |> Enum.to_list() + assert :rand.export_seed() == :undefined + end + + test "drop_remainder keeps only full batches" do + array = int_array(Enum.to_list(0..9)) + + sizes = + array + |> DataLoader.shuffled_batch_stream(4, seed: 1, drop_remainder: true) + |> Enum.map(fn {:ok, t} -> Nx.size(t) end) + + assert sizes == [4, 4] + end + end + describe "paired_batch_stream/4" do test "loads features and labels together" do {:ok, features} = create_test_array({100, 20}) diff --git a/test/ex_zarr/nx_test.exs b/test/ex_zarr/nx_test.exs index 8ee5679..3a826a4 100644 --- a/test/ex_zarr/nx_test.exs +++ b/test/ex_zarr/nx_test.exs @@ -236,6 +236,148 @@ defmodule ExZarr.NxTest do end end + describe "stream_chunk_tensors/2 options and edges" do + defp iota_array(shape, chunks, opts \\ []) do + {:ok, array} = + ExZarr.create([shape: shape, chunks: chunks, dtype: :float32, storage: :memory] ++ opts) + + start = List.to_tuple(List.duplicate(0, tuple_size(shape))) + data = Nx.iota(shape, type: {:f, 32}) |> Nx.to_binary() + :ok = ExZarr.Array.set_slice(array, data, start: start, stop: shape) + array + end + + # Every chunk tensor must equal the same region of to_tensor/2. + defp assert_chunks_match_whole(array, opts \\ []) do + {:ok, whole} = ExZarrNx.to_tensor(array) + + expected = + array + |> ExZarr.Array.stream_chunks(metadata: true) + |> Enum.map(fn %{index: index} -> + {start, stop} = ExZarr.Array.get_chunk_bounds(array, index) + starts = Tuple.to_list(start) + lengths = Enum.zip_with(starts, Tuple.to_list(stop), &(&2 - &1)) + whole |> Nx.slice(starts, lengths) |> Nx.to_flat_list() + end) + + actual = + array + |> ExZarrNx.stream_chunk_tensors(opts) + |> Enum.map(fn {:ok, tensor} -> Nx.to_flat_list(tensor) end) + + assert Enum.sort(actual) == Enum.sort(expected) + assert length(actual) == length(expected) + end + + test "crops edge chunks on every axis in 3D and 4D" do + assert_chunks_match_whole(iota_array({5, 7, 3}, {2, 3, 2})) + assert_chunks_match_whole(iota_array({3, 4, 5, 2}, {2, 3, 2, 2})) + assert_chunks_match_whole(iota_array({4, 6}, {4, 4})) + end + + test "unordered concurrent reads yield the right data for every chunk" do + assert_chunks_match_whole(iota_array({9, 9}, {2, 4}), concurrency: 4, ordered: false) + end + + test "include_missing yields fill-valued tensors cropped to the array bounds" do + {:ok, array} = + ExZarr.create( + shape: {3, 3}, + chunks: {2, 2}, + dtype: :int32, + fill_value: 7, + storage: :memory + ) + + :ok = ExZarr.Array.set_slice(array, <<1::32-signed-little>>, start: {0, 0}, stop: {1, 1}) + + stored = array |> ExZarrNx.stream_chunk_tensors() |> Enum.to_list() + assert [{:ok, only}] = stored + assert Nx.to_flat_list(only) == [1, 7, 7, 7] + + all = + array + |> ExZarrNx.stream_chunk_tensors(include_missing: true) + |> Enum.map(fn {:ok, tensor} -> {Nx.shape(tensor), Nx.to_flat_list(tensor)} end) + + assert all == [{{2, 2}, [1, 7, 7, 7]}, {{2, 1}, [7, 7]}, {{1, 2}, [7, 7]}, {{1, 1}, [7]}] + end + + test "an unsupported dtype yields one error and reads nothing" do + {:ok, array} = ExZarr.create(shape: {6}, chunks: {2}, dtype: :bool, storage: :memory) + + assert [{:error, message}] = + array |> ExZarrNx.stream_chunk_tensors(include_missing: true) |> Enum.to_list() + + assert message =~ "Unsupported dtype" + end + + test "applies the :backend option to each tensor" do + tensors = + {4} + |> iota_array({2}) + |> ExZarrNx.stream_chunk_tensors(backend: Nx.BinaryBackend) + |> Enum.map(fn {:ok, tensor} -> tensor end) + + assert Enum.all?(tensors, &match?(%Nx.Tensor{data: %Nx.BinaryBackend{}}, &1)) + end + + @tag :tmp_dir + test "on_error :skip drops an unreadable chunk and :halt raises", %{tmp_dir: tmp_dir} do + path = Path.join(tmp_dir, "corrupt") + + {:ok, array} = + ExZarr.create( + shape: {6}, + chunks: {2}, + dtype: :int32, + storage: :filesystem, + path: path, + compressor: :zlib + ) + + data = for i <- 0..5, into: <<>>, do: <> + :ok = ExZarr.Array.set_slice(array, data, start: {0}, stop: {6}) + + [chunk_file | _] = + path |> File.ls!() |> Enum.reject(&String.starts_with?(&1, ".")) |> Enum.sort() + + File.write!(Path.join(path, chunk_file), "not zlib data") + + skipped = array |> ExZarrNx.stream_chunk_tensors(on_error: :skip) |> Enum.to_list() + assert length(skipped) == 2 + + assert_raise ExZarr.StreamError, fn -> + array |> ExZarrNx.stream_chunk_tensors(on_error: :halt) |> Enum.to_list() + end + end + + @tag :performance + test "cropping a large padded edge chunk is fast" do + {:ok, array} = + ExZarr.create( + shape: {1999, 1000}, + chunks: {1000, 1000}, + dtype: :float32, + storage: :memory, + compressor: :none + ) + + data = :binary.copy(<<0::32>>, 1999 * 1000) + :ok = ExZarr.Array.set_slice(array, data, start: {0, 0}, stop: {1999, 1000}) + + {microseconds, [{:ok, tensor}]} = + :timer.tc(fn -> + array |> ExZarrNx.stream_chunk_tensors(filter: &(&1 == {1, 0})) |> Enum.to_list() + end) + + assert Nx.shape(tensor) == {999, 1000} + # The per-element crop in 1.3.0 took about 300 ms here; a row copy is ~1 ms. + assert microseconds < 100_000 + end + end + describe "to_tensor/2 - with options" do test "converts tensor with axis names" do {:ok, array} = diff --git a/test/mix/tasks/livebook_test_task_test.exs b/test/mix/tasks/livebook_test_task_test.exs new file mode 100644 index 0000000..2fed659 --- /dev/null +++ b/test/mix/tasks/livebook_test_task_test.exs @@ -0,0 +1,51 @@ +defmodule Mix.Tasks.Livebook.TestTaskTest do + use ExUnit.Case, async: true + + alias Mix.Tasks.Livebook.Test.LivemdExtract + + defp livemd(cells) do + Enum.map_join(cells, "\n\n", &("```elixir\n" <> &1 <> "\n```")) + end + + @tag :tmp_dir + test "cells run one at a time, so a later cell can use a struct defined at runtime", + %{tmp_dir: tmp_dir} do + # Compiled at runtime by cell 1, the way Mix.install/2 loads a package. + # Expanding the whole notebook up front fails on %LivebookCellStruct{}. + script = + LivemdExtract.to_elixir_script( + livemd([ + ~s|Code.compile_string("defmodule LivebookCellStruct do defstruct [:a] end")|, + ~s|alias LivebookCellStruct, as: S\nvalue = %S{a: 41}|, + ~s|%S{a: a} = value\nIO.puts("answer=" <> to_string(a + 1))| + ]) + ) + + path = Path.join(tmp_dir, "notebook.exs") + File.write!(path, script) + + assert {output, 0} = System.cmd("elixir", [path], stderr_to_stdout: true) + assert output =~ "answer=42" + end + + @tag :tmp_dir + test "a failing cell is named in the output", %{tmp_dir: tmp_dir} do + script = LivemdExtract.to_elixir_script(livemd(["x = 1", "raise \"boom\""])) + path = Path.join(tmp_dir, "failing.exs") + File.write!(path, script) + + assert {output, status} = System.cmd("elixir", [path], stderr_to_stdout: true) + assert status != 0 + assert output =~ "livebook cell 2 failed" + end + + test "--local points the Hex ex_zarr dependency at this checkout" do + cell = ~s|Mix.install([\n {:ex_zarr, "~> 1.3"},\n {:nx, "~> 0.7"}\n])| + root = File.cwd!() + + assert LivemdExtract.to_elixir_script(livemd([cell]), local: true) =~ + inspect(~s|{:ex_zarr, path: "#{root}"}|) |> String.slice(1..-2//1) + + assert LivemdExtract.to_elixir_script(livemd([cell])) =~ ~s|{:ex_zarr, \\"~> 1.3\\"}| + end +end