View Source Nx.LazyContainer protocol (Nx v0.5.0)
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
Functions
Traverses recursively tensors in a data structure with acc and fun.
Link to this section Types
@type t() :: term()
Link to this section Functions
@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.