A protocol that teaches Nx how to traverse data structures non-recursively.

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, ...]

Note the functions in this module are not recursive. If you want to deeply traverse and reduce containers, use the functions in Nx.Defn.Composite instead.

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.