View Source Nx.Container protocol (Nx v0.2.1)

A protocol that teaches Nx how to traverse data structures.

Nx and defn expect the arguments to be numbers, tensors, or one of the following composite data types:

  1. tuples of numbers/tensors
  2. maps of any key with numbers/tensors as values
  3. any struct that implements Nx.Container

If you need to pass additional values, you can implement or derive this protocol. For example:

@derive {Nx.Container,
         containers: [:field_name, :other_field]}
defstruct [:field_name, :other_fields, ...]

The :containers option is required and it must specify a list of fields that contains tensors. Inside defn, the container fields will be automatically converted to tensor expressions. All other fields will be reset to their default value, unless you explicitly declare them to be kept:

@derive {Nx.Container,
         containers: [:field_name, :other_field],
         keep: [:another_field]}
defstruct [:field_name, :other_fields, ...]

Careful!: If you keep a field, its value will be part of the Nx.Defn compiler cache key (i.e. therefore if you give a struct with two different values for a kept field, Nx.Defn will have to compile and cache it twice). You must only keep fields that you are certain to be used inside defn during compilation time.

Link to this section Summary

Functions

Reduces a data structure with acc and fun.

Traverse receives a data structure with acc and fun.

Link to this section Types

Link to this section Functions

@spec reduce(t(), acc, (Nx.Tensor.t(), acc -> acc)) :: acc when acc: term()

Reduces a data structure with acc and fun.

The function receives a tensor and the accumulator for each tensor in the container. It returns the update accumulator.

Link to this function

traverse(data, acc, fun)

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

Traverse receives a data structure with acc and fun.

The function receives a tensor and the accumulator for each tensor in the container. It returns a two element tuple with the updated container and the accumulator.