Nx
Nx v0.5.3

View Source Nx.LazyContainer protocol (Nx v0.5.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.

You can convert any Nx.LazyContainer to a tensor by using Nx.stack/2 and Nx.concatenate/2.

Link to this section 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.

Link to this section Types

Link to this type

t()

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

All the types that implement this protocol.

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.