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:
- tuples of numbers/tensors
- maps of any key with numbers/tensors as values
- 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 insidedefn
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
@type t() :: term()
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.
@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.