Machine Learning Pipelines

Mix.install([
  {:ex_zarr, "~> 1.3"},
  # {:ex_zarr, path: Path.join(__DIR__, "../../..")},
  {:nx, "~> 0.7"}
])

Streaming tensors from chunks

stream_chunk_tensors/2 reshapes each stored chunk, including a short edge chunk.

alias ExZarr.Array
alias ExZarr.Gallery.Pack

{:ok, array} =
  Array.create(
    shape: {100, 8},
    chunks: {32, 8},
    dtype: :float32,
    compressor: :zlib,
    storage: :memory
  )

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})

chunk_shapes =
  array
  |> ExZarr.Nx.stream_chunk_tensors(concurrency: 2, ordered: true)
  |> Enum.map(fn {:ok, tensor} -> Nx.shape(tensor) end)

%{chunk_shapes: chunk_shapes}

DataLoader (sample-level batches)

# Requires Nx; shuffled batches for training-style loops
array
|> ExZarr.Nx.DataLoader.shuffled_batch_stream(16, seed: 42)
|> Enum.take(2)
|> Enum.map(fn
  {:ok, batch} -> Nx.shape(batch)
  other -> other
end)

Broadway

For production prep, see docs/livebooks/broadway_pipeline.livemd and ExZarr.Broadway.start_chunk_pipeline/3.