Nx
Nx v0.5.1

View Source Nx.LazyContainer protocol (Nx v0.5.1)

Converts a data structure to a container lazily and recursively.

Sometimes building tensors for a container is an expensive operation, so we want to allow that to happen lazily.

This module provides a single traverse implementation that emits the tensor template and a function that computes the tensor as two distinct values. Then a tensor is only allocated if necessary.

This protocol is used throughout Nx.Defn API. This means compilation, jitting, and streaming will only realize lazy tensors when necessary.

If a data structure does not implement this protocol, a default implementation converts eager to lazy using Nx.Container. When a value is given to defn, it is first converted to tensors and containers via Nx.LazyContainer. Inside defn, there are no lazy containers, only containers.

Link to this section Summary

Types

t()

Functions

traverse(data, acc, fun)

Traverses recursively tensors in a data structure with acc and fun.

Link to this section Types

Link to this type

t()

View Source 
@type t() :: term()

Link to this section Functions

Link to this function

traverse(data, acc, fun)

View Source 
@spec traverse(t(), acc, (Nx.template(), (-> Nx.Tensor.t()), acc -> {term(), acc})) ::
  {Nx.Container.t(), acc}
when acc: term()

Traverses recursively tensors in a data structure with acc and fun.

For each tensor in the container, fun receives a tensor template, an anonymous function to build the actual tensor, and the accumulator . It returns a two element tuple with a non-lazy Nx.Container and the accumulator.

This function returns the updated container and the accumulator.

Note this function is recursive by default. Therefore if you are implementing this function and one of your arguments may be containers, you must call Nx.LazyContainer.traverse/3 on said arguments so they are recursively traversed.