The behaviour for tensor backends.
Each backend is module that defines a struct and implements the callbacks
defined in this module. The callbacks are mostly implementations of the
functions in the Nx module with the tensor output shape given as first
argument.
Nx backends come in two flavors: opaque backends, of which you should
not access its data directly except through the functions in the Nx
module, and public ones, of which its data can be directly accessed and
traversed. The former typically have the Backend suffix.
Nx ships with the following backends:
Nx.BinaryBackend- an opaque backend written in pure Elixir that stores the data in Elixir's binaries. This is the default backend used by theNxmodule. The backend itself (and its data) is private and must not be accessed directly.Nx.TemplateBackend- an opaque backend written that works as a template in APIs to declare the type, shape, and names of tensors to be expected in the future.Nx.Defn.Expr- a public backend used bydefnto build expression trees that are traversed by custom compilers.
This module also includes functions that are meant to be shared across backends.
Summary
Callbacks
Invoked for execution of Nx.block/4.
Functions
Inspects the given tensor given by binary.
Types
@type axes() :: Nx.Tensor.axes()
@type axis() :: Nx.Tensor.axis()
@type backend_options() :: term()
@type shape() :: Nx.Tensor.shape()
@type t() :: %{__struct__: atom()}
@type tensor() :: Nx.Tensor.t()
Callbacks
@callback backend_copy(tensor(), module(), backend_options()) :: tensor()
@callback backend_deallocate(tensor()) :: :ok | :already_deallocated
@callback backend_transfer(tensor(), module(), backend_options()) :: tensor()
@callback block(struct(), output :: tensor() | tuple(), args :: [term()], fun()) :: tensor() | tuple()
Invoked for execution of Nx.block/4.
output is the result template (Nx.Tensor or tuple of tensors). args are
the tensor (and optional trailing keyword lists) passed to Nx.block/4.
Backends should dispatch on struct (see Nx.Block.*) and either run a
native implementation or invoke fun as apply(fun, [struct | args]).
@callback constant(out :: tensor(), number() | Complex.t(), backend_options()) :: tensor()
@callback eye(tensor(), backend_options()) :: tensor()
@callback from_binary(out :: tensor(), binary(), backend_options()) :: tensor()
@callback init(keyword()) :: backend_options()
@callback inspect(tensor(), Inspect.Opts.t()) :: tensor()
@callback iota(tensor(), axis() | nil, backend_options()) :: tensor()
@callback to_binary(tensor(), limit :: non_neg_integer()) :: binary()