View Source Safetensors (Safetensors v0.1.3)

Safetensors implementation for Nx.

Examples

iex> x = Nx.tensor([1, 2, 3])
iex> y = Nx.tensor([1.0, 2.0, 3.0])
iex> tensors = %{"x" => x, "y" => y}
iex> data = Safetensors.dump(tensors)
iex> tensors = Safetensors.load!(data)
iex> tensors["x"]
#Nx.Tensor<
  s64[3]
  [1, 2, 3]
>
iex> tensors["y"]
#Nx.Tensor<
  f32[3]
  [1.0, 2.0, 3.0]
>

Summary

Functions

Serializes the given map of tensors to iodata.

Loads a serialized map of tensors.

Reads a serialized map of tensors from a file.

Writes a map of tensors to a file.

Functions

@spec dump(%{required(String.t()) => Nx.Tensor.t()}) :: iodata()

Serializes the given map of tensors to iodata.

iodata is a list of binaries that can be written to any IO device, such as a file or a socket. You can ensure the result is a binary by calling IO.iodata_to_binary/1.

@spec load!(iodata()) :: %{required(String.t()) => Nx.Tensor.t()}

Loads a serialized map of tensors.

It is the opposite of dump/1.

@spec read!(
  path :: Path.t(),
  keyword()
) :: %{required(String.t()) => Nx.LazyContainer.t()}

Reads a serialized map of tensors from a file.

Tensors are loaded into Nx one by one, without the need to load the entire file from disk into memory.

Options

  • :lazy - when true, instead of returning tensors, the function returns lazy containers. Such a container can be converted to a tensor using Nx.to_tensor/1 and it is only at that point that it is read from the file. Defaults to false
@spec write!(path :: Path.t(), %{required(String.t()) => Nx.Tensor.t()}) :: :ok

Writes a map of tensors to a file.

Tensors are written into the file one by one, without the need to dump all of them into the memory at once.