Powered by AppSignal & Oban Pro

Zarr to Nx

03_01_zarr_to_nx.livemd

Zarr to Nx

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.

alias ExZarr.Array

{:ok, array} =
  Array.create(shape: {6, 4}, chunks: {4, 2}, dtype: :float32, storage: :memory)

binary = for i <- 0..23, into: <<>>, do: <<i * 1.0::float-little-32>>
: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

{:ok, tensor} = ExZarr.Nx.to_tensor(array, names: [:rows, :cols])
{Nx.shape(tensor), Nx.type(tensor), Nx.names(tensor)}

Round trip

{: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.