From e20ce176e274a57c618eb8fb8e0320627f1fc940 Mon Sep 17 00:00:00 2001 From: thanos Date: Thu, 8 Oct 2026 14:32:00 -0400 Subject: [PATCH 1/3] ExZarr **1.3.0** is packaged as dense-array Nx interop. Arrow and Binsparse stay in ex_arrow and ex_graphblas. **Nx.** `ExZarr.Nx.stream_chunk_tensors/2` turns each stored chunk from `Array.stream_chunks/2` into a tensor. Edge chunks that are stored at the full chunk size are cropped to the array bounds. `:concurrency` and `:on_error` pass through. `DataLoader` docs now match the code: shuffled batches permute the full sample index, and `:shuffle_buffer_size` is ignored. 51 Nx tests passed, and `mix docs` built cleanly. **Docs and livebooks.** The Nx guide and `nx_streaming.livemd` use DataLoader and chunk streams. New notebooks cover metadata, Zarr-to-Nx conversion, minibatches, and a small Axon train. Mix.install cells in the embeddings and tick-data notebooks are fixed. The Axon notebook was not executed here because Axon is not a project dependency. **Release files.** `@version` is `1.3.0`, with changelog, release notes, roadmap, and README updated. The existing tag workflow already checks that `v1.3.0` matches `mix.exs` and `CHANGELOG.md`. I did not create or push that tag, so Hex is unchanged until you tag and push `v1.3.0`. - closed #310 - closed #311 - closed #312 - closed #313 --- CHANGELOG.md | 20 +- README.md | 8 +- docs/ROADMAP.md | 14 +- docs/guides/getting_started.md | 2 +- docs/guides/installation.md | 12 +- docs/guides/nx_integration.md | 165 +++++++--------- docs/guides/troubleshooting.md | 2 +- .../01_02_metadata_and_chunks.livemd | 111 +++++++++++ .../03_nx_ml/03_01_zarr_to_nx.livemd | 50 +++++ .../03_02_streaming_minibatches.livemd | 57 ++++++ .../03_nx_ml/03_03_training_from_zarr.livemd | 72 +++++++ .../04_01_embeddings_in_zarr.livemd | 6 +- .../05_finance/05_01_tick_data_cube.livemd | 6 +- .../06_cookbook/06_04_ml_pipelines.livemd | 20 +- docs/livebooks/PLAN.md | 15 ++ docs/livebooks/README.md | 4 + docs/livebooks/nx_streaming.livemd | 118 +++++++++--- docs/release_notes_v1_3_0.md | 51 +++++ lib/ex_zarr/nx.ex | 176 ++++++++++++++++-- lib/ex_zarr/nx/data_loader.ex | 27 ++- mix.exs | 12 +- test/ex_zarr/nx_test.exs | 71 +++++++ 22 files changed, 837 insertions(+), 182 deletions(-) create mode 100644 docs/livebooks/01_core_zarr/01_02_metadata_and_chunks.livemd create mode 100644 docs/livebooks/03_nx_ml/03_01_zarr_to_nx.livemd create mode 100644 docs/livebooks/03_nx_ml/03_02_streaming_minibatches.livemd create mode 100644 docs/livebooks/03_nx_ml/03_03_training_from_zarr.livemd create mode 100644 docs/release_notes_v1_3_0.md diff --git a/CHANGELOG.md b/CHANGELOG.md index 49d3b09..e063a75 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,23 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +## [1.3.0] - 2026-10-08 + +### Data Science Interop + +### Added +- `ExZarr.Nx.stream_chunk_tensors/2` maps `Array.stream_chunks/2` to Nx tensors. + Edge chunks stored at the full chunk size are cropped to the array bounds. + `:concurrency`, `:ordered`, `:timeout`, `:on_error`, `:filter`, and + `:include_missing` are forwarded. `:backend` and `:names` apply to each tensor. +- Livebooks: metadata and chunks, Zarr-to-Nx conversion, streaming minibatches, + and a small Axon training notebook. `nx_streaming.livemd` is the recipe book. + +### Changed +- `ExZarr.Nx.DataLoader` docs match the implementation: shuffled batches permute + the full sample index. `:shuffle_buffer_size` is accepted and ignored. +- Nx guide uses DataLoader and `stream_chunk_tensors/2` for batch and chunk streams. + ## [1.2.0] - 2026-10-04 ### Zarr 3.1 Interoperability & Range-Aware Cloud I/O @@ -885,7 +902,8 @@ None - Full backward compatibility maintained with v2 arrays --- -[Unreleased]: https://github.com/AnimaLogica/ExZarr/compare/v1.2.0...HEAD +[Unreleased]: https://github.com/AnimaLogica/ExZarr/compare/v1.3.0...HEAD +[1.3.0]: https://github.com/AnimaLogica/ExZarr/compare/v1.2.0...v1.3.0 [1.2.0]: https://github.com/AnimaLogica/ExZarr/releases/tag/v1.2.0 [1.1.0]: https://github.com/AnimaLogica/ExZarr/releases/tag/v1.1.0 [1.0.0]: https://github.com/AnimaLogica/ExZarr/releases/tag/v1.0.0 diff --git a/README.md b/README.md index 8118cb6..4c8cdfc 100644 --- a/README.md +++ b/README.md @@ -39,7 +39,7 @@ Add `ex_zarr` to your list of dependencies in `mix.exs`: ```elixir def deps do [ - {:ex_zarr, "~> 1.2"} + {:ex_zarr, "~> 1.3"} ] end ``` @@ -872,6 +872,10 @@ Key modules: See [ROADMAP.md](docs/ROADMAP.md) for the full release plan. +**v1.3.0 (released 2026-10-08)** - Data science interop: +`ExZarr.Nx.stream_chunk_tensors/2`, DataLoader batch recipes, livebook MVP +(metadata, Nx conversion, minibatches, a small Axon train). + **v1.2.0 (released 2026-10-04)** - Zarr 3.1 interoperability & range-aware cloud I/O: spec `sharding_indexed`, optional byte-range reads, version-aware cloud keys, AzureSDK migration, codecs via ExCodecs (no native toolchain), `compressor_config`, @@ -879,8 +883,6 @@ Python fixture/CI matrix. **Upcoming** (high level): -- **v1.3.0** - **Data science interop** - Explorer streaming, Nx batch recipes from `stream_chunks`, livebook curriculum, cookbook expansion. - **v1.4.0** - **Performance & packaging** Async codec pipeline (overlap I/O + decode), vendored/static codecs, adaptive range coalescing, storage-transformer execution. diff --git a/docs/ROADMAP.md b/docs/ROADMAP.md index 6354eb9..673bc69 100644 --- a/docs/ROADMAP.md +++ b/docs/ROADMAP.md @@ -26,11 +26,14 @@ Released 2026-06-12. Backward compatible with v1.0. --- -## v1.3.0 (Planned) — Data Science Interop +## v1.3.0 — Data Science Interop -- [ ] Explorer direct streaming integration -- [ ] `ExZarr.Nx` streaming recipes -- [ ] Livebook curriculum completion +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 --- @@ -55,4 +58,5 @@ Released 2026-06-12. Backward compatible with v1.0. - Unified retry layer across cloud SDKs - Native multi-range / coalescing heuristics - Partial writes inside cloud shards -- Explorer / new Nx APIs +- Partial shuffle buffer (`:shuffle_buffer_size` is accepted and ignored) +- Contiguous `load_samples_at_indices/3` reads diff --git a/docs/guides/getting_started.md b/docs/guides/getting_started.md index 6bd94ed..5f25f02 100644 --- a/docs/guides/getting_started.md +++ b/docs/guides/getting_started.md @@ -17,7 +17,7 @@ Add `ex_zarr` to your dependencies in `mix.exs`: ```elixir def deps do [ - {:ex_zarr, "~> 1.2"} + {:ex_zarr, "~> 1.3"} ] end ``` diff --git a/docs/guides/installation.md b/docs/guides/installation.md index a9cfe57..7e435aa 100644 --- a/docs/guides/installation.md +++ b/docs/guides/installation.md @@ -40,7 +40,7 @@ Add ExZarr to your `mix.exs` dependencies: ```elixir def deps do [ - {:ex_zarr, "~> 1.2"} + {:ex_zarr, "~> 1.3"} ] end ``` @@ -60,7 +60,7 @@ For cloud storage backends, add the appropriate libraries: ```elixir def deps do [ - {:ex_zarr, "~> 1.2"}, + {:ex_zarr, "~> 1.3"}, {:ex_aws, "~> 2.5"}, {:ex_aws_s3, "~> 2.5"}, {:sweet_xml, "~> 0.7"} # For XML response parsing @@ -72,7 +72,7 @@ end ```elixir def deps do [ - {:ex_zarr, "~> 1.2"}, + {:ex_zarr, "~> 1.3"}, {:goth, "~> 1.4"}, # Authentication {:req, "~> 0.6.1"} # HTTP client ] @@ -83,7 +83,7 @@ end ```elixir def deps do [ - {:ex_zarr, "~> 1.2"}, + {:ex_zarr, "~> 1.3"}, {:azure_sdk, "~> 0.4.1"} ] end @@ -93,7 +93,7 @@ end ```elixir def deps do [ - {:ex_zarr, "~> 1.2"}, + {:ex_zarr, "~> 1.3"}, {:mongodb_driver, "~> 1.4"} ] end @@ -181,7 +181,7 @@ If tests pass, your installation is complete and working correctly. ```elixir iex> Application.spec(:ex_zarr, :vsn) |> to_string() -"1.2.0" +"1.3.0" ``` ## Build Troubleshooting diff --git a/docs/guides/nx_integration.md b/docs/guides/nx_integration.md index 92a1e56..1b1dbb4 100644 --- a/docs/guides/nx_integration.md +++ b/docs/guides/nx_integration.md @@ -11,7 +11,7 @@ This guide shows how to integrate ExZarr with Nx (Numerical Elixir) for numerica - [Optimized Conversion (Recommended)](#optimized-conversion-recommended) - [Dtype and Shape Mapping](#dtype-and-shape-mapping) - [Chunked Processing](#chunked-processing) -- [Batch Processing Patterns](#batch-processing-patterns) +- [DataLoader and Chunk Streams](#dataloader-and-chunk-streams) - [ML Training/Inference Workflows](#ml-traininginference-workflows) - [Backend Integration](#backend-integration) - [Performance Considerations](#performance-considerations) @@ -23,7 +23,7 @@ This guide shows how to integrate ExZarr with Nx (Numerical Elixir) for numerica ```elixir # Install dependencies Mix.install([ - {:ex_zarr, "~> 1.2"}, + {:ex_zarr, "~> 1.3"}, {:nx, "~> 0.7"} ]) @@ -254,24 +254,40 @@ ExZarr.Nx.supported_nx_types() ## Chunked Processing -For arrays larger than available RAM, process in chunks with constant memory usage. +For arrays larger than available RAM, process one region at a time. -### Streaming Chunks as Tensors +### One tensor per stored chunk + +`ExZarr.Nx.stream_chunk_tensors/2` wraps `Array.stream_chunks/2`. Each item is the stored chunk, cropped to the array bounds when the last chunk is shorter than `chunks`. Pass `:concurrency` and `:on_error` through. ```elixir {:ok, array} = ExZarr.open(path: "/data/large_array") -# Process array in 100×100 chunks array -|> ExZarr.Nx.to_tensor_chunked({100, 100}) +|> ExZarr.Nx.stream_chunk_tensors(concurrency: 4) |> Stream.each(fn {:ok, tensor} -> - # Each tensor is 100×100 mean = Nx.mean(tensor) |> Nx.to_number() IO.puts("Chunk mean: #{mean}") end) |> Stream.run() ``` +### Regions that are not the stored chunk grid + +`to_tensor_chunked/3` reads rectangular slices of the size you pass. Use it when the region should not follow `array.chunks`. + +```elixir +{:ok, array} = ExZarr.open(path: "/data/large_array") + +array +|> ExZarr.Nx.to_tensor_chunked({100, 100}) +|> Stream.each(fn {:ok, tensor} -> + mean = Nx.mean(tensor) |> Nx.to_number() + IO.puts("Region mean: #{mean}") +end) +|> Stream.run() +``` + ### Parallel Chunk Processing ```elixir @@ -345,86 +361,51 @@ IO.puts("Global min: #{stats.min}") IO.puts("Global max: #{stats.max}") ``` -## Batch Processing Patterns - -### Mini-Batch Data Loading for ML - -Efficient batch loading for machine learning training: - -```elixir -defmodule BatchLoader do - @doc """ - Load batch from ExZarr array for ML training. - - Aligns ExZarr chunks with batch size for optimal performance. - """ - def load_batch(array, batch_idx, batch_size) do - {num_samples, num_features} = array.metadata.shape - - start_idx = batch_idx * batch_size - end_idx = min(start_idx + batch_size, num_samples) - - if start_idx >= num_samples do - {:error, :end_of_data} - else - # Load batch binary - {:ok, batch_binary} = ExZarr.Array.get_slice(array, - start: {start_idx, 0}, - stop: {end_idx, num_features} - ) - - # Convert to Nx tensor - {:ok, nx_type} = ExZarr.Nx.zarr_to_nx_type(array.metadata.dtype) +## DataLoader and Chunk Streams - batch_tensor = Nx.from_binary(batch_binary, nx_type) - |> Nx.reshape({end_idx - start_idx, num_features}) +Two streaming shapes cover the usual numerical workflows. - {:ok, batch_tensor} - end - end +| API | What each item is | Read path | +| --- | --- | --- | +| `ExZarr.Nx.stream_chunk_tensors/2` | One stored chunk | `Array.stream_chunks/2` | +| `ExZarr.Nx.DataLoader.batch_stream/3` | A fixed number of samples along axis 0 | `Array.get_slice/2` | +| `ExZarr.Nx.DataLoader.paired_batch_stream/4` | Matching feature and label batches | `get_slice/2` on both arrays | +| `ExZarr.Nx.to_tensor_chunked/3` | A rectangular region you choose | `get_slice/2` | - def batch_stream(array, batch_size) do - {num_samples, _} = array.metadata.shape - num_batches = ceil(num_samples / batch_size) +`batch_stream/3` can cross chunk boundaries. `stream_chunk_tensors/2` cannot: the tensor shape is the chunk that was stored, including a shorter edge chunk. - Stream.map(0..(num_batches - 1), fn batch_idx -> - load_batch(array, batch_idx, batch_size) - end) - |> Stream.take_while(fn - {:ok, _} -> true - {:error, :end_of_data} -> false - end) - end -end -``` +`shuffled_batch_stream/3` shuffles the full sample-index list for the epoch. `:shuffle_buffer_size` is accepted and ignored. -### Training Loop Example +### Training loop ```elixir -# Load training data from ExZarr -{:ok, X_train} = ExZarr.open(path: "/data/train_features") +{:ok, x_train} = ExZarr.open(path: "/data/train_features") {:ok, y_train} = ExZarr.open(path: "/data/train_labels") -# Training loop with batches model_state = initialize_model() -# Process batches -X_batches = BatchLoader.batch_stream(X_train, 32) -y_batches = BatchLoader.batch_stream(y_train, 32) +trained_state = + ExZarr.Nx.DataLoader.paired_shuffled_batch_stream(x_train, y_train, 32, seed: 1) + |> Enum.reduce(model_state, fn {:ok, {x_batch, y_batch}}, state -> + {loss, updated_state} = train_step(state, x_batch, y_batch) -trained_state = Stream.zip(X_batches, y_batches) -|> Enum.reduce(model_state, fn {{:ok, X_batch}, {:ok, y_batch}}, state -> - # Training step (with Nx.Defn for compilation) - {loss, updated_state} = train_step(state, X_batch, y_batch) + if rem(state.step, 100) == 0 do + IO.puts("Step #{state.step}, Loss: #{Float.round(loss, 4)}") + end - if rem(state.step, 100) == 0 do - IO.puts("Step #{state.step}, Loss: #{Float.round(loss, 4)}") - end + updated_state + end) +``` - updated_state -end) +### Chunk-aligned reduction + +```elixir +sums = + array + |> ExZarr.Nx.stream_chunk_tensors(concurrency: 4, on_error: :skip) + |> Enum.map(fn {:ok, tensor} -> Nx.sum(tensor) |> Nx.to_number() end) -IO.puts("Training complete!") +total = Enum.sum(sums) ``` ### Optimal Chunk Alignment @@ -617,15 +598,16 @@ defmodule Training do end # Load data outside defn (eager), compute inside (compiled) - def train_step(params, array, labels_array, batch_idx, batch_size) do - # Load batch from ExZarr (eager, not compiled) - {:ok, batch} = BatchLoader.load_batch(array, batch_idx, batch_size) - {:ok, labels} = BatchLoader.load_batch(labels_array, batch_idx, batch_size) - - # Compute with defn (lazy, compiled) - loss_value = loss(params, batch, labels) + def train_step(params, batch, labels) do + loss(params, batch, labels) + end - loss_value + def epoch(params, features, labels_array, batch_size) do + features + |> ExZarr.Nx.DataLoader.paired_batch_stream(labels_array, batch_size, drop_remainder: true) + |> Enum.reduce(params, fn {:ok, {batch, labels}}, params_acc -> + train_step(params_acc, batch, labels) + end) end end ``` @@ -717,11 +699,10 @@ compressor: %{id: "zstd", level: 10} # Best ratio # Avoid repeated allocations in loops defmodule EfficientProcessing do def process_batches(array, num_batches, batch_size) do - # Pre-allocate reusable tensor (if possible) - for batch_idx <- 0..(num_batches - 1) do - {:ok, batch} = BatchLoader.load_batch(array, batch_idx, batch_size) - process(batch) - end + array + |> ExZarr.Nx.DataLoader.batch_stream(batch_size, drop_remainder: true) + |> Stream.take(num_batches) + |> Enum.each(fn {:ok, batch} -> process(batch) end) end end ``` @@ -808,16 +789,10 @@ complex_restored = Nx.complex(real, imag) ```elixir # Load batches outside defn def train_epoch(params, array, batch_size) do - num_batches = div(elem(array.metadata.shape, 0), batch_size) - - Enum.reduce(0..(num_batches - 1), params, fn batch_idx, params_acc -> - # Load batch (eager, outside defn) - {:ok, batch} = BatchLoader.load_batch(array, batch_idx, batch_size) - - # Compute (lazy, compiled) - updated_params = train_step(params_acc, batch) - - updated_params + array + |> ExZarr.Nx.DataLoader.batch_stream(batch_size, drop_remainder: true) + |> Enum.reduce(params, fn {:ok, batch}, params_acc -> + train_step(params_acc, batch) end) end diff --git a/docs/guides/troubleshooting.md b/docs/guides/troubleshooting.md index ec7d73a..4d3a376 100644 --- a/docs/guides/troubleshooting.md +++ b/docs/guides/troubleshooting.md @@ -1244,7 +1244,7 @@ GitHub Issues: [https://github.com/AnimaLogica/ExZarr/issues](https://github.com ```markdown ## Environment -- ExZarr version: 1.2.0 +- ExZarr version: 1.3.0 - Elixir version: 1.19.5 - Erlang/OTP version: 28.1 - OS: macOS 15.2 (arm64) diff --git a/docs/livebooks/01_core_zarr/01_02_metadata_and_chunks.livemd b/docs/livebooks/01_core_zarr/01_02_metadata_and_chunks.livemd new file mode 100644 index 0000000..a4921fa --- /dev/null +++ b/docs/livebooks/01_core_zarr/01_02_metadata_and_chunks.livemd @@ -0,0 +1,111 @@ +# Metadata and chunks + +A Zarr array is metadata plus chunk files. This notebook writes v2 and v3 arrays, then reads the files ExZarr stored. + +```elixir +Mix.install([ + {:ex_zarr, "~> 1.3"}, + # {:ex_zarr, path: Path.join(__DIR__, "../../..")}, + {:kino, "~> 0.13"}, + {:jason, "~> 1.4"} +]) +``` + +## Setup + +```elixir +alias ExZarr.Array +alias ExZarr.Gallery.{Pack, SampleData} +``` + +## Create v2 and v3 arrays + +`:zarr_version` selects the metadata schema. The logical shape, chunks, and dtype stay the same. + +```elixir +common = [ + shape: {500, 400}, + chunks: {100, 80}, + dtype: :float32, + compressor: :zlib, + fill_value: 0.0, + storage: :memory +] + +{:ok, array_v2} = Array.create([zarr_version: 2] ++ common) +{:ok, array_v3} = Array.create([zarr_version: 3] ++ common) + +{array_v2.metadata.zarr_format, array_v3.metadata.zarr_format} +``` + +## Save and read metadata JSON + +```elixir +base_dir = Path.join(System.tmp_dir!(), "exzarr_metadata_demo") +v2_path = Path.join(base_dir, "array_v2") +v3_path = Path.join(base_dir, "array_v3") + +File.rm_rf!(base_dir) +:ok = Array.save(array_v2, path: v2_path) +:ok = Array.save(array_v3, path: v3_path) + +v2_json = v2_path |> Path.join(".zarray") |> File.read!() |> Jason.decode!() +v3_json = v3_path |> Path.join("zarr.json") |> File.read!() |> Jason.decode!() + +%{v2_keys: Map.keys(v2_json), v3_keys: Map.keys(v3_json)} +``` + +v2 stores a flat `.zarray` (`shape`, `chunks`, `dtype`, `compressor`). v3 stores `zarr.json` (`chunk_grid`, `data_type`, `codecs`, `chunk_key_encoding`). + +## Write a region and list chunk files + +A 150×150 write from the origin touches two row chunks and two column chunks. + +```elixir +binary = 150 |> SampleData.matrix(150) |> Pack.pack(:float32) +:ok = Array.set_slice(array_v2, binary, start: {0, 0}, stop: {150, 150}) +:ok = Array.set_slice(array_v3, binary, start: {0, 0}, stop: {150, 150}) +:ok = Array.save(array_v2, path: v2_path) +:ok = Array.save(array_v3, path: v3_path) + +v2_files = File.ls!(v2_path) |> Enum.sort() + +v3_files = + v3_path + |> Path.join("**/*") + |> Path.wildcard() + |> Enum.map(&Path.relative_to(&1, v3_path)) + |> Enum.sort() + +%{v2: v2_files, v3: v3_files} +``` + +v2 names chunks `0.0`, `0.1`, `1.0`, `1.1`. v3 nests them under `c/0/0`. + +## Chunk shape is an access-pattern choice + +```elixir +shapes = [{100, 100}, {50, 200}, {200, 50}] + +Enum.map(shapes, fn chunks -> + {:ok, array} = + Array.create(shape: {1000, 1000}, chunks: chunks, dtype: :int32, storage: :memory) + + {chunks, array.metadata.chunks} +end) +``` + +Square chunks suit mixed 2D reads. Wide chunks suit row scans. Tall chunks suit column scans. + +## Metadata on the struct + +```elixir +array_v2.metadata +``` + +## Takeaways + +1. Metadata is the contract; chunk files are the payload. +2. v2 uses `.zarray`. v3 uses `zarr.json` and a codec list. +3. Chunk filenames differ: `0.1` versus `c/0/1`. +4. Choose `chunks` for the reads you actually do. diff --git a/docs/livebooks/03_nx_ml/03_01_zarr_to_nx.livemd b/docs/livebooks/03_nx_ml/03_01_zarr_to_nx.livemd new file mode 100644 index 0000000..1d97ca6 --- /dev/null +++ b/docs/livebooks/03_nx_ml/03_01_zarr_to_nx.livemd @@ -0,0 +1,50 @@ +# Zarr to Nx + +Convert a whole array, a stored chunk, or a rectangular region into an `Nx.Tensor`. + +```elixir +Mix.install([ + {:ex_zarr, "~> 1.3"}, + # {:ex_zarr, path: Path.join(__DIR__, "../../..")}, + {:nx, "~> 0.7"} +]) +``` + +## Setup + +```elixir +alias ExZarr.Array + +{:ok, array} = + Array.create(shape: {6, 4}, chunks: {4, 2}, dtype: :float32, storage: :memory) + +binary = for i <- 0..23, into: <<>>, do: <> +:ok = Array.set_slice(array, binary, start: {0, 0}, stop: {6, 4}) + +{:ok, nx_type} = ExZarr.Nx.zarr_to_nx_type(array.metadata.dtype) +nx_type +``` + +## Whole array + +```elixir +{:ok, tensor} = ExZarr.Nx.to_tensor(array, names: [:rows, :cols]) +{Nx.shape(tensor), Nx.type(tensor), Nx.names(tensor)} +``` + +## Round trip + +```elixir +{:ok, stored} = + ExZarr.Nx.from_tensor(Nx.iota({8}, type: {:s, 32}), + storage: :memory, + chunks: {3} + ) + +{:ok, back} = ExZarr.Nx.to_tensor(stored) +Nx.to_flat_list(back) +``` + +## Dtypes that Zarr does not store + +`{:bf, 16}`, `{:f, 16}`, and complex types return `{:error, message}` from `nx_to_zarr_type/1`. Store float32 and cast in memory when you need those types. diff --git a/docs/livebooks/03_nx_ml/03_02_streaming_minibatches.livemd b/docs/livebooks/03_nx_ml/03_02_streaming_minibatches.livemd new file mode 100644 index 0000000..1d45a8d --- /dev/null +++ b/docs/livebooks/03_nx_ml/03_02_streaming_minibatches.livemd @@ -0,0 +1,57 @@ +# Streaming minibatches + +Chunk streams follow the stored grid. DataLoader batches follow the first axis and can cross chunk boundaries. + +```elixir +Mix.install([ + {:ex_zarr, "~> 1.3"}, + # {:ex_zarr, path: Path.join(__DIR__, "../../..")}, + {:nx, "~> 0.7"} +]) +``` + +## Data + +```elixir +alias ExZarr.Array + +{:ok, features} = + Array.create(shape: {10, 4}, chunks: {4, 4}, dtype: :float32, storage: :memory) + +{:ok, labels} = + Array.create(shape: {10}, chunks: {4}, dtype: :float32, storage: :memory) + +x = for i <- 0..39, into: <<>>, do: <> +y = for i <- 0..9, into: <<>>, do: <> +:ok = Array.set_slice(features, x, start: {0, 0}, stop: {10, 4}) +:ok = Array.set_slice(labels, y, start: {0}, stop: {10}) + +:ok +``` + +## One tensor per stored chunk + +The last chunk is shorter than `{4, 4}`. `stream_chunk_tensors/2` crops it. + +```elixir +features +|> ExZarr.Nx.stream_chunk_tensors(concurrency: 2, ordered: true) +|> Enum.map(fn {:ok, tensor} -> Nx.shape(tensor) end) +``` + +## Fixed-size batches + +```elixir +features +|> ExZarr.Nx.DataLoader.batch_stream(3, drop_remainder: true) +|> Enum.map(fn {:ok, batch} -> Nx.shape(batch) end) +``` + +## Paired and shuffled batches + +Shuffling uses the full sample index. `:shuffle_buffer_size` is ignored. `:seed` makes the order repeatable. + +```elixir +ExZarr.Nx.DataLoader.paired_shuffled_batch_stream(features, labels, 4, seed: 1) +|> Enum.map(fn {:ok, {xb, yb}} -> {Nx.shape(xb), Nx.shape(yb)} end) +``` 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 new file mode 100644 index 0000000..6a1a28a --- /dev/null +++ b/docs/livebooks/03_nx_ml/03_03_training_from_zarr.livemd @@ -0,0 +1,72 @@ +# Training from Zarr + +A one-layer Axon model trained on batches read from ExZarr. Axon is installed only in this notebook. + +```elixir +Mix.install([ + {:ex_zarr, "~> 1.3"}, + # {:ex_zarr, path: Path.join(__DIR__, "../../..")}, + {:nx, "~> 0.7"}, + {:axon, "~> 0.7"} +]) +``` + +## Store a linear dataset + +Targets are `2 * x0 + 1` plus a little noise stored as float32. + +```elixir +alias ExZarr.Array + +n = 64 +features = for(i <- 0..(n - 1), j <- 0..1, into: <<>>, do: <<(i + j) * 0.05::float-little-32>>) + +targets = + for i <- 0..(n - 1), into: <<>> do + <<(2.0 * (i * 0.05) + 1.0)::float-little-32>> + end + +{:ok, x_array} = + Array.create(shape: {n, 2}, chunks: {16, 2}, dtype: :float32, storage: :memory) + +{:ok, y_array} = + Array.create(shape: {n, 1}, chunks: {16, 1}, dtype: :float32, storage: :memory) + +: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 +``` + +## Model and batches + +```elixir +model = + Axon.input("features", shape: {nil, 2}) + |> Axon.dense(1) + +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) + +model +``` + +## A few SGD steps + +```elixir +trainer = Axon.Loop.trainer(model, :mean_squared_error, :sgd) + +trained = Axon.Loop.run(trainer, batches, %{}, epochs: 3, iterations: 4) + +model_state = + case trained do + %Axon.Loop.State{step_state: %{model_state: state}} -> state + state -> state + end + +Axon.predict(model, model_state, %{"features" => Nx.tensor([[0.0, 0.0], [1.0, 0.0]])}) +``` diff --git a/docs/livebooks/04_ai_genai/04_01_embeddings_in_zarr.livemd b/docs/livebooks/04_ai_genai/04_01_embeddings_in_zarr.livemd index 9e1198b..0d197df 100644 --- a/docs/livebooks/04_ai_genai/04_01_embeddings_in_zarr.livemd +++ b/docs/livebooks/04_ai_genai/04_01_embeddings_in_zarr.livemd @@ -1,8 +1,10 @@ # Embeddings in Zarr (chunked vectors + similarity scan) ```elixir - {:ex_zarr, "~> 1.2"}, - # {:ex_zarr, path: Path.join(__DIR__, "../../..")}, +Mix.install([ + {:ex_zarr, "~> 1.3"}, + # {:ex_zarr, path: Path.join(__DIR__, "../../..")} +]) ``` ## Section diff --git a/docs/livebooks/05_finance/05_01_tick_data_cube.livemd b/docs/livebooks/05_finance/05_01_tick_data_cube.livemd index 80ea18c..10e986b 100644 --- a/docs/livebooks/05_finance/05_01_tick_data_cube.livemd +++ b/docs/livebooks/05_finance/05_01_tick_data_cube.livemd @@ -1,8 +1,10 @@ # Finance: tick data cube (time × symbol × field) ```elixir - {:ex_zarr, "~> 1.2"}, - # {:ex_zarr, path: Path.join(__DIR__, "../../..")}, +Mix.install([ + {:ex_zarr, "~> 1.3"}, + # {:ex_zarr, path: Path.join(__DIR__, "../../..")} +]) ``` ## Section diff --git a/docs/livebooks/06_cookbook/06_04_ml_pipelines.livemd b/docs/livebooks/06_cookbook/06_04_ml_pipelines.livemd index 5c9f81e..19dc4b9 100644 --- a/docs/livebooks/06_cookbook/06_04_ml_pipelines.livemd +++ b/docs/livebooks/06_cookbook/06_04_ml_pipelines.livemd @@ -2,7 +2,7 @@ ```elixir Mix.install([ - {:ex_zarr, "~> 1.2"}, + {:ex_zarr, "~> 1.3"}, # {:ex_zarr, path: Path.join(__DIR__, "../../..")}, {:nx, "~> 0.7"} ]) @@ -10,30 +10,30 @@ Mix.install([ ## Streaming tensors from chunks +`stream_chunk_tensors/2` reshapes each stored chunk, including a short edge chunk. + ```elixir alias ExZarr.Array alias ExZarr.Gallery.Pack {:ok, array} = Array.create( - shape: {128, 8}, + shape: {100, 8}, chunks: {32, 8}, dtype: :float32, compressor: :zlib, storage: :memory ) -vals = for i <- 0..(128 * 8 - 1), do: i * 0.01 -:ok = Array.set_slice(array, Pack.pack(vals, :float32), start: {0, 0}, stop: {128, 8}) +vals = for i <- 0..(100 * 8 - 1), do: i * 0.01 +:ok = Array.set_slice(array, Pack.pack(vals, :float32), start: {0, 0}, stop: {100, 8}) -batches = +chunk_shapes = array - |> Array.stream_chunks(concurrency: 2) - |> Stream.map(fn {_index, data} -> Nx.from_binary(data, {:f, 32}) end) - |> Stream.chunk_every(2) - |> Enum.to_list() + |> ExZarr.Nx.stream_chunk_tensors(concurrency: 2, ordered: true) + |> Enum.map(fn {:ok, tensor} -> Nx.shape(tensor) end) -%{batch_count: length(batches), first_batch_shapes: Enum.map(hd(batches), &Nx.shape/1)} +%{chunk_shapes: chunk_shapes} ``` ## DataLoader (sample-level batches) diff --git a/docs/livebooks/PLAN.md b/docs/livebooks/PLAN.md index 9439225..516854e 100644 --- a/docs/livebooks/PLAN.md +++ b/docs/livebooks/PLAN.md @@ -1,3 +1,18 @@ +## Status (v1.3.0) + +Shipped for 1.3.0: + +- `01_core_zarr/01_01` through `01_04` +- `03_nx_ml/03_01_zarr_to_nx.livemd` +- `03_nx_ml/03_02_streaming_minibatches.livemd` +- `03_nx_ml/03_03_training_from_zarr.livemd` +- `04_ai_genai/04_01_embeddings_in_zarr.livemd` +- `05_finance/05_01_tick_data_cube.livemd` +- `06_cookbook/06_01` through `06_07` +- `nx_streaming.livemd` + +Deferred past 1.3.0: `02_concurrency`, `04_02`–`04_15`, `05_02`–`05_12`, `06_crypto` through `09_systems`, and the rest of the table below. Arrow and Explorer livebooks belong with ex_arrow. Binsparse livebooks belong with ex_graphblas. + The livebooks directry should be a curated collection of **Elixir Livebooks** demonstrating how to use **ExZarr** — a Zarr v3–compatible array storage library — to build scalable data systems for **AI / GenAI, finance, crypto, and scientific workloads**. diff --git a/docs/livebooks/README.md b/docs/livebooks/README.md index 61ebba1..1c1dad9 100644 --- a/docs/livebooks/README.md +++ b/docs/livebooks/README.md @@ -26,7 +26,11 @@ Or open any `.livemd` in [Livebook](https://livebook.dev/). Path installs use | Livebook | Topic | |----------|--------| | [`01_core_zarr/01_01_first_zarr_array.livemd`](01_core_zarr/01_01_first_zarr_array.livemd) | Create → write → read → stream → save | +| [`01_core_zarr/01_02_metadata_and_chunks.livemd`](01_core_zarr/01_02_metadata_and_chunks.livemd) | v2/v3 metadata and chunk files | | [`01_core_zarr/01_03_chunk_streaming.livemd`](01_core_zarr/01_03_chunk_streaming.livemd) | Sequential vs parallel chunk streaming | +| [`03_nx_ml/03_01_zarr_to_nx.livemd`](03_nx_ml/03_01_zarr_to_nx.livemd) | Array, tensor, and dtype conversion | +| [`03_nx_ml/03_02_streaming_minibatches.livemd`](03_nx_ml/03_02_streaming_minibatches.livemd) | Chunk streams and DataLoader batches | +| [`03_nx_ml/03_03_training_from_zarr.livemd`](03_nx_ml/03_03_training_from_zarr.livemd) | Small Axon train from Zarr batches | | [`01_core_zarr/01_04_codecs_and_pipelines.livemd`](01_core_zarr/01_04_codecs_and_pipelines.livemd) | Codecs and pipelines | | [`04_ai_genai/04_01_embeddings_in_zarr.livemd`](04_ai_genai/04_01_embeddings_in_zarr.livemd) | Embeddings in Zarr | | [`05_finance/05_01_tick_data_cube.livemd`](05_finance/05_01_tick_data_cube.livemd) | Tick data cube | diff --git a/docs/livebooks/nx_streaming.livemd b/docs/livebooks/nx_streaming.livemd index 8eb54d9..1b74d42 100644 --- a/docs/livebooks/nx_streaming.livemd +++ b/docs/livebooks/nx_streaming.livemd @@ -2,70 +2,130 @@ ## Introduction -Stream Zarr chunks directly into Nx tensors for machine learning workflows. +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`) ```elixir Mix.install([ - {:ex_zarr, "~> 1.2"}, - # {:ex_zarr, path: Path.join(__DIR__, "../../..")} + {:ex_zarr, "~> 1.3"}, + # {:ex_zarr, path: Path.join(__DIR__, "../..")} {:nx, "~> 0.7"}, {:kino, "~> 0.13"} ]) ``` -## Create Array and Stream to Tensors +## Section + +Create a small in-memory array. The last chunk is shorter than `chunks`. ```elixir alias ExZarr.Array {:ok, array} = Array.create( - shape: {100, 100}, - chunks: {10, 10}, + shape: {10, 4}, + chunks: {4, 4}, dtype: :float32, storage: :memory ) -data = for(_ <- 1..100, into: <<>>, do: <<1.0::float-little-32>>) -Array.set_slice(array, data, start: {0, 0}, stop: {10, 10}) +data = for i <- 0..39, into: <<>>, do: <> +:ok = Array.set_slice(array, data, start: {0, 0}, stop: {10, 4}) + +:ok +``` + +## Section + +`stream_chunk_tensors/2` yields `{:ok, tensor}`. Edge chunks are cropped to the array bounds. `:concurrency` and `:on_error` are forwarded to `stream_chunks/2`. + +```elixir +chunk_shapes = + array + |> ExZarr.Nx.stream_chunk_tensors(concurrency: 2, ordered: true) + |> Enum.map(fn {:ok, tensor} -> Nx.shape(tensor) end) + +Kino.inspect(chunk_shapes) +``` + +## Section -tensors = +Sum each stored chunk, then add the partial sums. This stays chunk-aligned and does not build a sample index. + +```elixir +total = array - |> Array.stream_chunks() - |> Stream.take(4) - |> Enum.map(fn {_index, chunk_data} -> - Nx.from_binary(chunk_data, {:f, 32}) + |> ExZarr.Nx.stream_chunk_tensors() + |> Enum.reduce(0.0, fn {:ok, tensor}, acc -> + acc + (Nx.sum(tensor) |> Nx.to_number()) end) -length(tensors) +total +``` + +## Section + +`to_tensor_chunked/3` ignores the stored grid and reads regions of the size you pass. + +```elixir +region_shapes = + array + |> ExZarr.Nx.to_tensor_chunked({3, 4}) + |> Enum.map(fn {:ok, tensor} -> Nx.shape(tensor) end) + +region_shapes +``` + +## Section + +`DataLoader.batch_stream/3` slices along axis 0. A batch can cross chunk boundaries. `drop_remainder: true` drops a short final batch. + +```elixir +batch_shapes = + array + |> ExZarr.Nx.DataLoader.batch_stream(3, drop_remainder: true) + |> Enum.map(fn {:ok, batch} -> Nx.shape(batch) end) + +batch_shapes ``` -## Batch Pipeline +## Section + +Shuffling permutes every sample index for the epoch. `:shuffle_buffer_size` is ignored. Pass `:seed` when you need the same order again. ```elixir -batches = +first_rows = array - |> Array.stream_chunks(concurrency: 2) - |> Stream.map(fn {_index, data} -> Nx.from_binary(data, {:f, 32}) end) - |> Stream.chunk_every(2) + |> ExZarr.Nx.DataLoader.shuffled_batch_stream(2, seed: 7, drop_remainder: true) |> Enum.take(2) + |> Enum.map(fn {:ok, batch} -> Nx.to_list(batch) end) -Kino.inspect(batches |> Enum.map(fn batch -> Enum.map(batch, &Nx.shape/1) end)) +first_rows ``` -## Using DataLoader +## Section -For training workflows with shuffling: +Pair a feature array with a label array. Both must have the same length on axis 0. ```elixir -array -|> ExZarr.Nx.DataLoader.batch_stream(32) -|> Enum.take(3) -|> Enum.map(fn {:ok, batch} -> Nx.shape(batch) end) +{:ok, labels} = + Array.create(shape: {10}, chunks: {4}, dtype: :float32, storage: :memory) + +label_data = for i <- 0..9, into: <<>>, do: <<(i * 10.0)::float-little-32>> +:ok = Array.set_slice(labels, label_data, start: {0}, stop: {10}) + +pairs = + ExZarr.Nx.DataLoader.paired_batch_stream(array, labels, 4) + |> Enum.map(fn {:ok, {x, y}} -> {Nx.shape(x), Nx.shape(y)} end) + +pairs ``` ## Key Takeaways -1. `stream_chunks` composes naturally with `Nx.from_binary/2` -2. Use `concurrency` to overlap I/O with tensor construction -3. `ExZarr.Nx.DataLoader` handles sample-level batching and shuffling +1. Use `stream_chunk_tensors/2` for chunk-aligned map/reduce and `:concurrency`. +2. Use `DataLoader.batch_stream/3` when the step needs a fixed sample count. +3. Use `paired_batch_stream/4` so features and labels stay aligned. +4. `shuffled_batch_stream/3` keeps the full index list in memory for the epoch. diff --git a/docs/release_notes_v1_3_0.md b/docs/release_notes_v1_3_0.md new file mode 100644 index 0000000..dc9a237 --- /dev/null +++ b/docs/release_notes_v1_3_0.md @@ -0,0 +1,51 @@ +# v1.3.0 Release Notes + +## Data Science Interop + +ExZarr v1.3.0 adds Nx recipes for stored chunks and sample batches, plus a +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. + +## Highlights + +### Chunk tensors + +`ExZarr.Nx.stream_chunk_tensors/2` wraps `Array.stream_chunks/2`. Each binary +becomes an `Nx.Tensor` using the array dtype. When a chunk is stored at the +full chunk size and the array ends inside it, the tensor is the in-bounds +region. + +```elixir +array +|> ExZarr.Nx.stream_chunk_tensors(concurrency: 4) +|> 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. + +### Batches + +`ExZarr.Nx.DataLoader.batch_stream/3` and `paired_batch_stream/4` slice along +the first axis with `get_slice/2`. A batch can cross chunk boundaries. + +`shuffled_batch_stream/3` shuffles every sample index for the epoch. +`:shuffle_buffer_size` is accepted and ignored. Pass `:seed` for a repeatable +order. + +### Livebooks + +- `docs/livebooks/01_core_zarr/01_02_metadata_and_chunks.livemd` +- `docs/livebooks/03_nx_ml/03_01_zarr_to_nx.livemd` +- `docs/livebooks/03_nx_ml/03_02_streaming_minibatches.livemd` +- `docs/livebooks/03_nx_ml/03_03_training_from_zarr.livemd` +- `docs/livebooks/nx_streaming.livemd` + +```elixir +{:ex_zarr, "~> 1.3"} +``` + +Hex owner remains `thanos`. Source: https://github.com/AnimaLogica/ExZarr diff --git a/lib/ex_zarr/nx.ex b/lib/ex_zarr/nx.ex index 3bbe30e..9b8421d 100644 --- a/lib/ex_zarr/nx.ex +++ b/lib/ex_zarr/nx.ex @@ -42,7 +42,8 @@ defmodule ExZarr.Nx do {:ok, tensors} = ExZarr.Nx.to_tensor_chunked(array, {100, 100}) # Returns stream of smaller tensors - See `ExZarr.Nx.DataLoader` for ML training workflows. + See `ExZarr.Nx.DataLoader` for sample-axis batches. Use `stream_chunk_tensors/2` + when you want one tensor per stored chunk via `Array.stream_chunks/2`. """ alias ExZarr.Array @@ -240,45 +241,81 @@ defmodule ExZarr.Nx do end @doc """ - Converts ExZarr array to stream of Nx tensors by processing in chunks. + Streams one Nx tensor per Zarr chunk. - For large arrays that don't fit in memory, this function loads the array - in chunks and yields a tensor for each chunk. Useful for processing large - datasets incrementally. + Wraps `ExZarr.Array.stream_chunks/2`. Each yielded chunk binary is reshaped + with that chunk's bounds (edge chunks are smaller than `array.chunks`) and + the array dtype. Streaming options are forwarded unchanged. + + Use this when work is chunk-aligned (decode, reduce, write-back). Use + `ExZarr.Nx.DataLoader` when you need fixed-size batches along the first + axis; those read with `get_slice/2` and can span chunk boundaries. + + ## Options + + Tensor options (not forwarded to `stream_chunks/2`): + + - `:backend` - Nx backend to transfer each tensor to + - `:names` - Axis names for each tensor + + All other options match `ExZarr.Array.stream_chunks/2` (`:concurrency`, + `:ordered`, `:timeout`, `:on_error`, `:filter`, `:include_missing`). + `: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`. + + ## Examples + + array + |> ExZarr.Nx.stream_chunk_tensors(concurrency: 4) + |> Stream.map(fn {:ok, tensor} -> Nx.mean(tensor) end) + |> Enum.to_list() + + """ + @spec stream_chunk_tensors(ExZarr.Array.t(), keyword()) :: + 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)) + end + + @doc """ + Converts ExZarr array to stream of Nx tensors by processing in slices. + + For large arrays that don't fit in memory, this function loads rectangular + regions of `chunk_size` and yields a tensor for each region. Regions follow + the requested size, not the array's stored chunk grid. For one tensor per + stored chunk, use `stream_chunk_tensors/2`. ## Parameters - `array` - ExZarr array to convert - - `chunk_size` - Size of chunks to load (tuple matching array dimensionality) + - `chunk_size` - Size of regions to load (tuple matching array dimensionality) - `opts` - Options (same as `to_tensor/2`) ## Returns - Stream that yields `{:ok, tensor}` or `{:error, reason}` for each chunk. + Stream that yields `{:ok, tensor}` or `{:error, reason}` for each region. ## Examples - # Process 100×100 chunks from large array + # Process 100×100 regions from a large array {:ok, array} = ExZarr.open(path: "/data/large_array") array |> ExZarr.Nx.to_tensor_chunked({100, 100}) |> Stream.each(fn {:ok, tensor} -> - # Process each tensor chunk result = Nx.mean(tensor) |> Nx.to_number() - IO.puts("Chunk mean: \#{result}") + IO.puts("Region mean: \#{result}") end) |> Stream.run() - # Map over chunks in parallel - results = - array - |> ExZarr.Nx.to_tensor_chunked({100, 100}) - |> Task.async_stream(fn {:ok, tensor} -> - Nx.sum(tensor) |> Nx.to_number() - end, max_concurrency: 4) - |> Enum.to_list() - """ @spec to_tensor_chunked(ExZarr.Array.t(), tuple(), keyword()) :: Enumerable.t({:ok, Nx.Tensor.t()} | {:error, term()}) @@ -475,4 +512,105 @@ defmodule ExZarr.Nx do |> Enum.map(fn {start_val, stop_val} -> stop_val - start_val end) |> List.to_tuple() end + + defp chunk_event_to_tensor(array, %{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 + tensor = + payload + |> Nx.from_binary(nx_type) + |> Nx.reshape(shape) + |> apply_tensor_opts(opts) + + {:ok, tensor} + end + end + + defp crop_chunk_binary(data, shape, nominal, itemsize) do + logical_bytes = tuple_product(shape) * itemsize + nominal_bytes = tuple_product(nominal) * itemsize + + cond do + byte_size(data) == logical_bytes -> + {:ok, data} + + byte_size(data) == nominal_bytes -> + {:ok, extract_logical_chunk(data, shape, nominal, itemsize)} + + true -> + {:error, + {:invalid_chunk_size, + %{expected: logical_bytes, nominal: nominal_bytes, actual: byte_size(data)}}} + end + end + + defp extract_logical_chunk(data, shape, nominal, itemsize) do + strides = chunk_strides(nominal) + + shape + |> tuple_coordinates() + |> Enum.map(fn coord -> + offset = flat_offset(coord, strides) * itemsize + binary_part(data, offset, itemsize) + 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() + end + + defp flat_offset(coord, strides) do + Enum.zip(coord, strides) + |> Enum.reduce(0, fn {index, stride}, acc -> acc + index * stride end) + end + + defp tuple_product(shape) do + shape |> Tuple.to_list() |> Enum.reduce(1, &*/2) + end + + defp apply_tensor_opts(tensor, opts) do + tensor = + case Keyword.get(opts, :names) do + nil -> tensor + names -> Nx.rename(tensor, names) + end + + case Keyword.get(opts, :backend) do + nil -> tensor + backend -> Nx.backend_transfer(tensor, backend) + end + end end diff --git a/lib/ex_zarr/nx/data_loader.ex b/lib/ex_zarr/nx/data_loader.ex index 6731d0c..bc6fe95 100644 --- a/lib/ex_zarr/nx/data_loader.ex +++ b/lib/ex_zarr/nx/data_loader.ex @@ -56,6 +56,20 @@ defmodule ExZarr.Nx.DataLoader do |> Enum.each(fn {:ok, {x_batch, y_batch}} -> train_step(model, x_batch, y_batch) end) + + ## Batches vs chunk streams + + `batch_stream/3` and `shuffled_batch_stream/3` slice along the first axis with + `ExZarr.Array.get_slice/2`. A batch can cross stored chunk boundaries, and + the batch size is independent of `array.chunks`. + + `ExZarr.Nx.stream_chunk_tensors/2` yields one tensor per stored chunk and + forwards `:concurrency` to `stream_chunks/2`. Use it for chunk-aligned + reductions. Use this module when the training step needs a fixed sample count. + + `shuffled_batch_stream/3` materializes and shuffles every sample index for + the epoch. `:shuffle_buffer_size` is accepted for compatibility and ignored; + there is no partial shuffle buffer. """ alias ExZarr.Array @@ -66,7 +80,7 @@ defmodule ExZarr.Nx.DataLoader do - `:drop_remainder` - Drop incomplete final batch (default: false) - `:seed` - Random seed for shuffling (default: nil, uses random) - - `:shuffle_buffer_size` - Size of shuffle buffer (default: batch_size * 10) + - `:shuffle_buffer_size` - Accepted and ignored. Shuffling uses the full index list. - `:backend` - Nx backend to transfer tensors to (default: nil) - `:names` - Axis names for tensors (default: nil) """ @@ -156,19 +170,18 @@ defmodule ExZarr.Nx.DataLoader do Streams shuffled batches from ExZarr array for ML training. Randomizes sample order before batching, which typically improves training - convergence. Uses a shuffle buffer to balance memory usage and randomization quality. + convergence. ## Shuffling Strategy - Uses reservoir sampling with a configurable buffer size: - - Larger buffer: Better randomization, more memory - - Smaller buffer: Less memory, local randomization - - Default: 10x batch size + 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 + O(number of samples). `:shuffle_buffer_size` is ignored. ## Options - `:seed` - Random seed for reproducibility (default: random) - - `:shuffle_buffer_size` - Buffer size for shuffling (default: batch_size * 10) + - `:shuffle_buffer_size` - Accepted and ignored - `:drop_remainder` - Drop incomplete final batch (default: false) - `:backend` - Nx backend to transfer tensors to (default: nil) - `:names` - Axis names for tensors (default: nil) diff --git a/mix.exs b/mix.exs index a2f1142..c6a9fe7 100644 --- a/mix.exs +++ b/mix.exs @@ -1,7 +1,7 @@ defmodule ExZarr.MixProject do use Mix.Project - @version "1.2.0" + @version "1.3.0" @source_url "https://github.com/AnimaLogica/ExZarr" def project do @@ -163,7 +163,11 @@ defmodule ExZarr.MixProject do # Examples "docs/livebooks/README.md", "docs/livebooks/01_core_zarr/01_01_first_zarr_array.livemd", + "docs/livebooks/01_core_zarr/01_02_metadata_and_chunks.livemd", "docs/livebooks/01_core_zarr/01_03_chunk_streaming.livemd", + "docs/livebooks/03_nx_ml/03_01_zarr_to_nx.livemd", + "docs/livebooks/03_nx_ml/03_02_streaming_minibatches.livemd", + "docs/livebooks/03_nx_ml/03_03_training_from_zarr.livemd", "docs/livebooks/01_core_zarr/01_04_codecs_and_pipelines.livemd", "docs/livebooks/04_ai_genai/04_01_embeddings_in_zarr.livemd", "docs/livebooks/05_finance/05_01_tick_data_cube.livemd", @@ -194,6 +198,7 @@ defmodule ExZarr.MixProject do "CHANGELOG.md", "docs/release_notes_v1_1_0.md", "docs/release_notes_v1_2_0.md", + "docs/release_notes_v1_3_0.md", "docs/ROADMAP.md", "docs/ZARR_V3_STATUS.md", "docs/INTEROPERABILITY.md", @@ -237,7 +242,11 @@ defmodule ExZarr.MixProject do Examples: [ "docs/livebooks/README.md", "docs/livebooks/01_core_zarr/01_01_first_zarr_array.livemd", + "docs/livebooks/01_core_zarr/01_02_metadata_and_chunks.livemd", "docs/livebooks/01_core_zarr/01_03_chunk_streaming.livemd", + "docs/livebooks/03_nx_ml/03_01_zarr_to_nx.livemd", + "docs/livebooks/03_nx_ml/03_02_streaming_minibatches.livemd", + "docs/livebooks/03_nx_ml/03_03_training_from_zarr.livemd", "docs/livebooks/01_core_zarr/01_04_codecs_and_pipelines.livemd", "docs/livebooks/04_ai_genai/04_01_embeddings_in_zarr.livemd", "docs/livebooks/05_finance/05_01_tick_data_cube.livemd", @@ -278,6 +287,7 @@ defmodule ExZarr.MixProject do "Release Notes & Migration Guides": [ "docs/release_notes_v1_1_0.md", "docs/release_notes_v1_2_0.md", + "docs/release_notes_v1_3_0.md", "docs/migration_guide_v1_1_0.md" ], Reference: [ diff --git a/test/ex_zarr/nx_test.exs b/test/ex_zarr/nx_test.exs index 9cf2d1e..8ee5679 100644 --- a/test/ex_zarr/nx_test.exs +++ b/test/ex_zarr/nx_test.exs @@ -165,6 +165,77 @@ defmodule ExZarr.NxTest do end end + describe "stream_chunk_tensors/2" do + test "yields one tensor per stored chunk, including a short edge chunk" do + {:ok, array} = + ExZarr.create( + shape: {5}, + chunks: {3}, + dtype: :int32, + storage: :memory + ) + + data = for i <- 0..4, into: <<>>, do: <> + :ok = ExZarr.Array.set_slice(array, data, start: {0}, stop: {5}) + + tensors = + array + |> ExZarrNx.stream_chunk_tensors(names: [:n]) + |> Enum.map(fn {:ok, tensor} -> tensor end) + + assert length(tensors) == 2 + assert Enum.map(tensors, &Nx.shape/1) == [{3}, {2}] + assert Enum.all?(tensors, &(Nx.names(&1) == [:n])) + assert tensors |> Enum.map(&Nx.to_flat_list/1) |> List.flatten() == [0, 1, 2, 3, 4] + end + + test "crops a padded edge chunk in C order" do + {:ok, array} = + ExZarr.create( + shape: {2, 3}, + chunks: {2, 2}, + dtype: :uint8, + storage: :memory + ) + + data = <<1, 2, 3, 4, 5, 6>> + :ok = ExZarr.Array.set_slice(array, data, start: {0, 0}, stop: {2, 3}) + + values = + array + |> ExZarrNx.stream_chunk_tensors(ordered: true) + |> Enum.map(fn {:ok, tensor} -> Nx.to_flat_list(tensor) end) + + assert values == [[1, 2, 4, 5], [3, 6]] + end + + test "forwards concurrency without treating it as a tensor option" do + {:ok, array} = + ExZarr.create( + shape: {4, 4}, + chunks: {2, 2}, + dtype: :uint8, + storage: :memory + ) + + data = for i <- 0..15, into: <<>>, do: <> + :ok = ExZarr.Array.set_slice(array, data, start: {0, 0}, stop: {4, 4}) + + results = + array |> ExZarrNx.stream_chunk_tensors(concurrency: 2, ordered: true) |> Enum.to_list() + + assert length(results) == 4 + assert Enum.all?(results, &match?({:ok, _}, &1)) + + assert Enum.map(results, fn {:ok, tensor} -> Nx.shape(tensor) end) == [ + {2, 2}, + {2, 2}, + {2, 2}, + {2, 2} + ] + end + end + describe "to_tensor/2 - with options" do test "converts tensor with axis names" do {:ok, array} = From 3b0e8e9f534c7782825dffda6ffb6d1dde97fa66 Mon Sep 17 00:00:00 2001 From: thanos Date: Thu, 8 Oct 2026 16:41:59 -0400 Subject: [PATCH 2/3] Nx is the Elixir library for doing the same math across a whole grid of numbers at once, instead of looping in Elixir one element at a time. You use it for: - **Summaries** over an image, a table, or a cube of measurements: sum, mean, min, max - **Linear algebra**: dot products and matrix multiplies - **Training**: the prediction step and the gradient of a loss. Axon builds models on top of Nx - **Compiled runs**: `defn` turns those expressions into CPU code, or GPU code through EXLA or Torchx The grid stays in memory as a tensor: one binary buffer, a type such as `{:f, 32}`, and a shape. ExZarr does not run that math. It stores the bytes. `ExZarr.Nx` copies an array into a tensor when you want to compute, and copies a tensor back when you want to keep the result. The getting-started section **What is Nx** now says this directly, without assuming you know NumPy. `03_03_training_from_zarr.livemd` now explains each step before the code runs. * **The dataset** is 64 rows of two features. The target is `2 * x0 + 1`. Column 1 is stored and does not enter that formula. Chunks of 16 match the batch size, so each step reads one chunk. The cell prints the first row `[0.0, 0.05]` and the first target `1.0`. * **The model** is one dense layer: `y_hat = w0 * x0 + w1 * x1 + b`. A good fit is weights near `[2, 0]` and bias near `1`. The batch cell shows shapes `{16, 2}` and `{16, 1}`, and how that pair is wrapped as `%{"features" => x}` for Axon. * **Training** is mean squared error and SGD for 3 passes of 4 batches, 12 updates from random weights. The two probe rows are `[0.0, 0.0]` and `[1.0, 0.0]`, whose line values are `1` and `3`. The notebook says those predictions move toward `1` and `3` and will not land on them after this short run. The first target is stored as one column, so `y0[0]` has shape `{1}` and `Nx.to_number/1` rejects it. The cell now indexes that remaining column with `y0[0][0]` before converting to a number. --- README.md | 4 +- docs/guides/getting_started.md | 84 +++++++++++++++++++ docs/guides/nx_integration.md | 9 +- .../03_nx_ml/03_01_zarr_to_nx.livemd | 4 +- .../03_02_streaming_minibatches.livemd | 38 +++++++-- .../03_nx_ml/03_03_training_from_zarr.livemd | 72 ++++++++++++++-- docs/livebooks/scripts/00_01_welcome.exs | 13 ++- examples/nx_integration.exs | 4 +- 8 files changed, 193 insertions(+), 35 deletions(-) diff --git a/README.md b/README.md index 4c8cdfc..62c7a47 100644 --- a/README.md +++ b/README.md @@ -21,6 +21,7 @@ interop tested in CI against zarr-python 2.18 / 3.2 / 3.3 / 3.4. See - **High Performance** - 26x faster multi-chunk reads with near-optimal scaling (see [Performance Guide](docs/guides/performance.md)) - **N-dimensional arrays** with support for 10 data types (int8-64, uint8-64, float32/64) - **BEAM-native streaming** - `stream_chunks/2`, `stream_slices/3`, and `write_stream/3` for bounded-memory processing +- **Nx** - `ExZarr.Nx` converts arrays and chunks to tensors; `ExZarr.Nx.DataLoader` streams sample-axis batches for training. Nx is optional. See [Using Nx](docs/guides/getting_started.md#using-nx) - **Telemetry** - `:telemetry` events for chunk I/O and stream lifecycle (`ExZarr.Telemetry`) - **Pipeline integrations** - Optional Flow, GenStage, and Broadway support for production pipelines - **Parallel chunk processing** - Automatic parallel I/O and decompression for large operations @@ -832,7 +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 -- **[Nx Integration](examples/nx_integration.exs)** - Numerical computing with Nx +- **[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/guides/getting_started.md b/docs/guides/getting_started.md index 5f25f02..abc5842 100644 --- a/docs/guides/getting_started.md +++ b/docs/guides/getting_started.md @@ -429,6 +429,89 @@ Common errors: - **`:not_found`**: Array or chunk not found - **`:dimension_mismatch`**: Data shape doesn't match slice dimensions +## Using Nx + +Nx is optional. Add `{:nx, "~> 0.7"}` next to `ex_zarr` when you want in-memory numerical work on data that lives in a Zarr array. + +### What is Nx + +[Nx](https://hexdocs.pm/nx) is the Elixir library for math on a grid of numbers. You use it when the same operation should hit every element, row, or matrix at once, instead of walking the data in an Elixir loop. + +Typical jobs: + +- Summaries: sum, mean, min, max over an image, a table, or a cube of measurements +- Linear algebra: dot products and matrix multiplies +- Training: the prediction step and the gradient of a loss. [Axon](https://hexdocs.pm/axon) builds models on top of Nx +- Compiled runs: `defn` turns those expressions into code for the CPU, or for a GPU through EXLA or Torchx + +A tensor is that grid in memory: one contiguous binary, a type such as `{:f, 32}`, a shape tuple, and a backend. The default backend is `Nx.BinaryBackend`. + +ExZarr does not run that math. It keeps the bytes on disk or in a store so the process can exit and so another language can open the same array. `ExZarr.Nx` copies a Zarr array into a tensor when you want to compute, and copies a tensor back when you want to store the result. + +### When to Use It + +Use `ExZarr.Nx` when the numbers are already in a Zarr array, or when a tensor should be written as one: + +- The array is larger than you want in RAM, so you stream chunks or batches. +- Training data and labels should stay aligned on the first axis. +- A result tensor should be saved as Zarr and reopened from Python. +- Model weights or an intermediate cube should be checkpointed as an array. + +```elixir +{:ok, array} = ExZarr.open(path: "/data/features") +{:ok, tensor} = ExZarr.Nx.to_tensor(array) +Nx.mean(tensor) +``` + +### When not to use it + +Skip the bridge when the tensor already fits in memory and nothing else needs to read it. Stay on `Nx` alone. + +Skip it when the access pattern is a few scalar lookups. A Zarr read loads a chunk, so random single-element access pays I/O and decompression for a whole chunk. + +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 + +| Function | What you get | +| --- | --- | +| `ExZarr.Nx.to_tensor/2` | The whole array as one tensor | +| `ExZarr.Nx.from_tensor/2` | A new array filled from a tensor | +| `ExZarr.Nx.stream_chunk_tensors/2` | One tensor per stored chunk. The last chunk is cropped to the array bounds | +| `ExZarr.Nx.to_tensor_chunked/3` | Tensors for rectangular regions you choose, which can differ from `chunks` | +| `ExZarr.Nx.DataLoader.batch_stream/3` | Batches of a fixed number of samples along axis 0 | +| `ExZarr.Nx.DataLoader.paired_batch_stream/4` | Matching feature and label batches | +| `ExZarr.Nx.zarr_to_nx_type/1` | `:float32` → `{:f, 32}`, and the other nine numeric dtypes | + +`stream_chunk_tensors/2` forwards `:concurrency` and `:on_error` to `Array.stream_chunks/2`. `:names` and `:backend` apply to each tensor. + +`DataLoader` reads with `get_slice/2`, so a batch can cross a chunk boundary. `shuffled_batch_stream/3` permutes every sample index for the epoch. `:seed` repeats that order. `:shuffle_buffer_size` is accepted and ignored. + +```elixir +array +|> ExZarr.Nx.stream_chunk_tensors(concurrency: 4) +|> Enum.map(fn {:ok, chunk} -> Nx.sum(chunk) end) +``` + +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 + +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. + +A Zarr array is a directory (or an object prefix, or an in-memory map). `zarr.json` or `.zarray` records shape, chunk shape, and dtype. The numbers sit in separate chunk objects. An array of shape `{6, 4}` with chunks `{4, 4}` is three chunk slots along the first axis: two full `{4, 4}` chunks and one edge chunk whose logical shape is `{2, 4}`. ExZarr often stores that edge chunk padded out to the full chunk shape, with the fill value past the array bounds. + +The element types match on the ten numeric dtypes. `:float32` on disk is `{:f, 32}` in Nx. Both sides use little-endian IEEE layout for the values ExZarr writes. The names differ (`:float32` versus `{:f, 32}`); `zarr_to_nx_type/1` and `nx_to_zarr_type/1` translate them. Shapes are tuples on both sides, so `{6, 4}` does not need a reshape of the *logical* array, only of each chunk binary. + +Exchange is a copy through that binary layout: + +1. **Array to tensor.** `to_tensor/2` reads every chunk, decompresses it, and concatenates the bytes in C order into one binary. `Nx.from_binary/2` builds the tensor and `Nx.reshape/2` attaches the array shape. The tensor then lives in the backend you asked for (`:backend` calls `Nx.backend_transfer/2`). +2. **Tensor to array.** `from_tensor/2` reads `Nx.to_binary/1`, creates an array with the tensor shape and the dtype you mapped, and writes chunk-sized slices. You choose `chunks`. That choice is about later reads, not about how Nx laid the tensor out. +3. **One chunk at a time.** `stream_chunk_tensors/2` does not concatenate. Each chunk binary becomes its own tensor. If the stored bytes are the padded full chunk, the helper keeps only the in-bounds region before `Nx.reshape/2`. +4. **A batch of samples.** `DataLoader` calls `get_slice/2` for a range on axis 0. That read may open two chunks, stitch the requested rows, and hand `Nx.from_binary/2` a binary whose length is `batch * rest_of_shape * itemsize`. + +There is no shared memory between the array and the tensor. Decompression, the slice stitch, and `from_binary` each produce a new binary. After the tensor exists, Nx operations do not touch the Zarr store until you write again with `from_tensor/2` or `set_slice/3`. + ## Next Steps Now that you understand the basics: @@ -444,6 +527,7 @@ Now that you understand the basics: - `sharded_cloud_storage.exs`: Using sharding with S3 - `dimension_names.exs`: Named dimension slicing - `nx_integration.exs`: Integration with Nx numerical library + - The [Using Nx](#using-nx) section above, and [Nx integration](nx_integration.md) for the full API 3. **Read API documentation**: Browse the full API at https://hexdocs.pm/ex_zarr diff --git a/docs/guides/nx_integration.md b/docs/guides/nx_integration.md index 1b1dbb4..5e8c7a2 100644 --- a/docs/guides/nx_integration.md +++ b/docs/guides/nx_integration.md @@ -2,6 +2,8 @@ This guide shows how to integrate ExZarr with Nx (Numerical Elixir) for numerical computing and machine learning workflows. ExZarr provides persistent storage for Nx tensors, enabling workflows that exceed available memory. +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. ## Table of Contents @@ -46,12 +48,7 @@ tensor = Nx.iota({1000, 1000}) ### The Relationship -**Nx** (Numerical Elixir) provides NumPy-like functionality for Elixir: -- In-memory numerical computing -- Multi-dimensional tensors -- Vectorized operations -- Defn numerical definitions -- Backend support (CPU, GPU via EXLA) +**Nx** is how you run math on a whole grid of numbers in Elixir: sums and means, matrix multiplies, and the gradients a training step needs. The grid stays in memory as a tensor. `defn` compiles those expressions for the CPU, or for a GPU through EXLA or Torchx. [Axon](https://hexdocs.pm/axon) builds models on top of Nx. See [What is Nx](getting_started.md#what-is-nx) for a longer description. **ExZarr** provides persistent, chunked array storage: - Disk/cloud-backed arrays diff --git a/docs/livebooks/03_nx_ml/03_01_zarr_to_nx.livemd b/docs/livebooks/03_nx_ml/03_01_zarr_to_nx.livemd index 1d97ca6..dcb414b 100644 --- a/docs/livebooks/03_nx_ml/03_01_zarr_to_nx.livemd +++ b/docs/livebooks/03_nx_ml/03_01_zarr_to_nx.livemd @@ -1,7 +1,5 @@ # Zarr to Nx -Convert a whole array, a stored chunk, or a rectangular region into an `Nx.Tensor`. - ```elixir Mix.install([ {:ex_zarr, "~> 1.3"}, @@ -12,6 +10,8 @@ Mix.install([ ## Setup +Convert a whole array, a stored chunk, or a rectangular region into an `Nx.Tensor`. + ```elixir alias ExZarr.Array diff --git a/docs/livebooks/03_nx_ml/03_02_streaming_minibatches.livemd b/docs/livebooks/03_nx_ml/03_02_streaming_minibatches.livemd index 1d45a8d..ae4f03a 100644 --- a/docs/livebooks/03_nx_ml/03_02_streaming_minibatches.livemd +++ b/docs/livebooks/03_nx_ml/03_02_streaming_minibatches.livemd @@ -1,7 +1,5 @@ # Streaming minibatches -Chunk streams follow the stored grid. DataLoader batches follow the first axis and can cross chunk boundaries. - ```elixir Mix.install([ {:ex_zarr, "~> 1.3"}, @@ -10,7 +8,19 @@ Mix.install([ ]) ``` -## Data +## Setup + +This notebook builds a tiny feature matrix and a label vector, then reads them three ways. + +* One tensor per **stored chunk**. The grid is whatever you set in `chunks`. +* One tensor per **minibatch**. You choose how many samples (axis 0) each step sees. A batch can cross a chunk boundary. +* **Paired** feature and label batches, with a shuffle that is repeatable when you pass `:seed`. + +## A 10×4 feature matrix and 10 labels + +`features` has shape `{10, 4}` and chunks `{4, 4}`. Along the sample axis that is three chunks: rows 0–3, rows 4–7, and rows 8–9. The last chunk is shorter than 4 rows. + +`labels` is a length-10 vector with chunks of 4, so it has the same sample boundaries. Values are just the index, stored as little-endian float32, which is the layout `Nx.from_binary/2` expects for `{:f, 32}`. ```elixir alias ExZarr.Array @@ -26,12 +36,16 @@ y = for i <- 0..9, into: <<>>, do: <> :ok = Array.set_slice(features, x, start: {0, 0}, stop: {10, 4}) :ok = Array.set_slice(labels, y, start: {0}, stop: {10}) -:ok +{features.metadata.shape, features.chunks, labels.metadata.shape, labels.chunks} ``` ## One tensor per stored chunk -The last chunk is shorter than `{4, 4}`. `stream_chunk_tensors/2` crops it. +`stream_chunk_tensors/2` follows the chunk grid. It does not invent a batch size. + +`:concurrency` is forwarded to `Array.stream_chunks/2` (here 2 readers). `:ordered` keeps chunk index order, so the shapes come out as `{4, 4}`, `{4, 4}`, then the edge `{2, 4}`. The edge chunk is stored padded to `{4, 4}`; the helper crops it to the two rows that belong to the array. + +Use this path for reductions and maps that can run per chunk. The tensor shape changes on the last chunk, so a training step that requires a static batch size should use DataLoader instead. ```elixir features @@ -39,7 +53,11 @@ features |> Enum.map(fn {:ok, tensor} -> Nx.shape(tensor) end) ``` -## Fixed-size batches +## Minibatches of 3 samples + +`batch_stream/3` slices axis 0 with `get_slice/2`. Each item is `{:ok, tensor}` of shape `{3, 4}`: 3 samples, all 4 features. + +`drop_remainder: true` drops the final incomplete group. 10 samples ÷ 3 leaves 1 sample, so you get three batches and the last sample is omitted. The second batch starts at row 3, which is inside the first chunk, and ends at row 6, which is inside the second chunk. That single batch reads across the chunk boundary. A chunk stream cannot do that. ```elixir features @@ -47,9 +65,13 @@ features |> Enum.map(fn {:ok, batch} -> Nx.shape(batch) end) ``` -## Paired and shuffled batches +## Shuffled feature/label pairs + +`paired_shuffled_batch_stream/4` draws the same sample indexes for `features` and `labels`, so row `i` of a feature batch is still label `i`. + +The shuffle builds the full index list `0..9`, permutes it, then loads each batch. `:seed` makes that permutation repeatable. `:shuffle_buffer_size` is accepted and ignored. There is no partial buffer. -Shuffling uses the full sample index. `:shuffle_buffer_size` is ignored. `:seed` makes the order repeatable. +Batch size 4 with 10 samples yields shapes `{4, 4}` / `{4}`, then `{4, 4}` / `{4}`, then `{2, 4}` / `{2}`. The short last pair is the remainder. Pass `drop_remainder: true` when the training step needs every batch to have the same size. ```elixir ExZarr.Nx.DataLoader.paired_shuffled_batch_stream(features, labels, 4, seed: 1) 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 6a1a28a..c33cb89 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 @@ -1,19 +1,32 @@ # Training from Zarr -A one-layer Axon model trained on batches read from ExZarr. Axon is installed only in this notebook. - ```elixir Mix.install([ {:ex_zarr, "~> 1.3"}, # {:ex_zarr, path: Path.join(__DIR__, "../../..")}, {:nx, "~> 0.7"}, - {:axon, "~> 0.7"} + {:axon, "~> 0.7"} # Axon is installed only in this notebook. ExZarr does not depend on it. ]) ``` -## Store a linear dataset +## Intro + +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. + +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. + +## Write the line into two arrays + +There are 64 samples. Each feature row is two float32 numbers: -Targets are `2 * x0 + 1` plus a little noise stored as float32. +* column 0: `i * 0.05` +* column 1: `(i + 1) * 0.05` + +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`. + +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. ```elixir alias ExZarr.Array @@ -35,10 +48,28 @@ 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 +{:ok, x0} = ExZarr.Nx.to_tensor(x_array) +{:ok, y0} = ExZarr.Nx.to_tensor(y_array) + +%{ + 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() +} ``` -## Model and batches +The map should show `x_shape: {64, 2}`, `y_shape: {64, 1}`, `first_row: [0.0, 0.05]`, and `first_target: 1.0`. + +## 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. + +`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"`. `Nx.as_type/2` makes the dtype explicit before the training loop compiles. ```elixir model = @@ -52,10 +83,31 @@ batches = end) |> Stream.map(fn {x, y} -> {%{"features" => x}, y} end) -model +{first_inputs, first_targets} = batches |> Enum.take(1) |> hd() + +%{ + model: model, + batch_x: Nx.shape(first_inputs["features"]), + batch_y: Nx.shape(first_targets) +} ``` -## A few SGD steps +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. + +`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.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. + +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) @@ -70,3 +122,5 @@ model_state = Axon.predict(model, model_state, %{"features" => Nx.tensor([[0.0, 0.0], [1.0, 0.0]])}) ``` + +The result is a `{2, 1}` tensor. Row 0 is the prediction for `x0 = 0`. Row 1 is the prediction for `x0 = 1`. diff --git a/docs/livebooks/scripts/00_01_welcome.exs b/docs/livebooks/scripts/00_01_welcome.exs index 4a2fce6..a5cd2e8 100644 --- a/docs/livebooks/scripts/00_01_welcome.exs +++ b/docs/livebooks/scripts/00_01_welcome.exs @@ -175,17 +175,16 @@ Vl.new(width: 300, height: 300, title: "Zarr Array: Divided into 16 Chunks (4×4 # ### Why Nx? Understanding Elixir's Numerical Computing Library -# Nx (Numerical Elixir) is Elixir's numerical computing library, similar to NumPy in Python. ExZarr integrates deeply with Nx for several compelling reasons: +# Nx is the Elixir library for math on a grid of numbers: sums and means, matrix multiplies, and the gradients a training step needs. ExZarr stores those grids; Nx computes on them. # #### What is Nx? -# Nx provides multi-dimensional tensors (arrays) with: +# You use Nx when the same operation should hit every element of a multi-dimensional array at once. A tensor is that grid in memory. -# * **Type safety:** Compile-time shape and type checking -# * **Defn compilation:** JIT compilation of numerical functions to CPU/GPU -# * **Backend flexibility:** CPU (BinaryBackend), GPU (EXLA/Torchx), TPU support -# * **Lazy evaluation:** Build computation graphs before execution -# * **Broadcasting:** Automatic shape alignment like NumPy +# * **Summaries and linear algebra:** sum, mean, dot products, matrix multiplies +# * **Training:** prediction steps and gradients. Axon builds models on Nx +# * **Defn compilation:** those expressions compile for the CPU, or for a GPU through EXLA or Torchx +# * **Backends:** CPU (`BinaryBackend`) by default; EXLA and Torchx for accelerators # #### Why ExZarr Uses Nx diff --git a/examples/nx_integration.exs b/examples/nx_integration.exs index 5407fd4..95e8fb1 100755 --- a/examples/nx_integration.exs +++ b/examples/nx_integration.exs @@ -19,8 +19,8 @@ defmodule NxIntegrationExample do @moduledoc """ Demonstrates integration between ExZarr and Nx. - Nx provides NumPy-like functionality for Elixir, making it a natural - fit for numerical computing with Zarr arrays. + Nx runs math on a whole grid of numbers in memory: summaries, matrix + multiplies, and training steps. ExZarr stores those grids as Zarr arrays. """ def run do From 5d35cf6f2599f7373e5d9b2c12c68956c7c2df25 Mon Sep 17 00:00:00 2001 From: thanos Date: Thu, 8 Oct 2026 18:10:37 -0400 Subject: [PATCH 3/3] `mix docs --warnings-as-errors` completes cleanly now. ExDoc only links to pages listed in `mix.exs`. The Using Nx links pointed at `docs/guides/getting_started.md`, and that file was not registered, so the anchors were reported as missing files. It is now included in the docs extras and in the Getting Started group. --- mix.exs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/mix.exs b/mix.exs index c6a9fe7..cb25e2c 100644 --- a/mix.exs +++ b/mix.exs @@ -140,6 +140,7 @@ defmodule ExZarr.MixProject do extras: [ # Getting Started "docs/guides/what_is_zarr.md", + "docs/guides/getting_started.md", {"README.md", title: "Introduction"}, # Core Concepts @@ -219,6 +220,7 @@ defmodule ExZarr.MixProject do groups_for_extras: [ "Getting Started": [ "docs/guides/what_is_zarr.md", + "docs/guides/getting_started.md", "README.md" ], "Core Concepts": [