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..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 @@ -39,7 +40,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 ``` @@ -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 @@ -872,6 +874,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 +885,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..abc5842 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 ``` @@ -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/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..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 @@ -11,7 +13,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 +25,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"} ]) @@ -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 @@ -254,24 +251,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. + +### One tensor per stored chunk -### Streaming Chunks as Tensors +`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 +358,51 @@ IO.puts("Global min: #{stats.min}") IO.puts("Global max: #{stats.max}") ``` -## Batch Processing Patterns +## DataLoader and Chunk Streams -### Mini-Batch Data Loading for ML +Two streaming shapes cover the usual numerical workflows. -Efficient batch loading for machine learning training: +| 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` | -```elixir -defmodule BatchLoader do - @doc """ - Load batch from ExZarr array for ML training. +`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. - 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) - - batch_tensor = Nx.from_binary(batch_binary, nx_type) - |> Nx.reshape({end_idx - start_idx, num_features}) +`shuffled_batch_stream/3` shuffles the full sample-index list for the epoch. `:shuffle_buffer_size` is accepted and ignored. - {:ok, batch_tensor} - end - end - - def batch_stream(array, batch_size) do - {num_samples, _} = array.metadata.shape - num_batches = ceil(num_samples / batch_size) - - 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 -``` - -### 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 +595,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 +696,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 +786,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..dcb414b --- /dev/null +++ b/docs/livebooks/03_nx_ml/03_01_zarr_to_nx.livemd @@ -0,0 +1,50 @@ +# Zarr to Nx + +```elixir +Mix.install([ + {:ex_zarr, "~> 1.3"}, + # {:ex_zarr, path: Path.join(__DIR__, "../../..")}, + {:nx, "~> 0.7"} +]) +``` + +## Setup + +Convert a whole array, a stored chunk, or a rectangular region into an `Nx.Tensor`. + +```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..ae4f03a --- /dev/null +++ b/docs/livebooks/03_nx_ml/03_02_streaming_minibatches.livemd @@ -0,0 +1,79 @@ +# Streaming minibatches + +```elixir +Mix.install([ + {:ex_zarr, "~> 1.3"}, + # {:ex_zarr, path: Path.join(__DIR__, "../../..")}, + {:nx, "~> 0.7"} +]) +``` + +## 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 + +{: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}) + +{features.metadata.shape, features.chunks, labels.metadata.shape, labels.chunks} +``` + +## One tensor per stored chunk + +`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 +|> ExZarr.Nx.stream_chunk_tensors(concurrency: 2, ordered: true) +|> Enum.map(fn {:ok, tensor} -> Nx.shape(tensor) end) +``` + +## 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 +|> ExZarr.Nx.DataLoader.batch_stream(3, drop_remainder: true) +|> Enum.map(fn {:ok, batch} -> Nx.shape(batch) end) +``` + +## 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. + +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) +|> 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..c33cb89 --- /dev/null +++ b/docs/livebooks/03_nx_ml/03_03_training_from_zarr.livemd @@ -0,0 +1,126 @@ +# Training from Zarr + +```elixir +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. +]) +``` + +## 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: + +* 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 + +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, 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() +} +``` + +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 = + 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) + +{first_inputs, first_targets} = batches |> Enum.take(1) |> hd() + +%{ + model: model, + batch_x: Nx.shape(first_inputs["features"]), + batch_y: Nx.shape(first_targets) +} +``` + +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) + +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]])}) +``` + +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/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/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/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/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 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..cb25e2c 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 @@ -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 @@ -163,7 +164,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 +199,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", @@ -214,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": [ @@ -237,7 +244,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 +289,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} =