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

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

Types

t()

All the types that implement this protocol.

Functions

traverse(data, acc, fun)

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

Types

t()

@type t() :: term()

All the types that implement this protocol.

Functions

traverse(data, acc, fun)

@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.