View Source Nx.LazyContainer protocol (Nx v0.6.1)
Converts a data structure to a container lazily.
When you invoke a defn, its arguments must implement
the Nx.LazyContainer protocol and return a data structure
that implements Nx.Container. In other words, this
protocol teaches Nx to work with additional data types
beyond numbers and tensors.
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.
Nx.LazyContainer is automatically provided for data
structures that implement Nx.Container.
Summary
Functions
Traverses recursively tensors in a data structure with acc and fun.
Types
@type t() :: term()
All the types that implement this protocol.
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.