View Source Nx.Backend behaviour (Nx v0.6.0)
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 theNx
module. 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 bydefn
to 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 optional callbacks with a default implementation.
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 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()
Invoked for execution of optional callbacks with a default implementation.
First we will attempt to call the optional callback itself (one of the many callbacks defined below), then we attempt to call this callback (which is also optional), then we fallback to the default iomplementation.
@callback random_normal(tensor(), mu :: tensor(), sigma :: tensor(), backend_options()) :: tensor()
@callback random_uniform(tensor(), tensor(), tensor(), backend_options()) :: tensor()
@callback to_binary(tensor(), limit :: non_neg_integer()) :: binary()
Functions
Inspects the given tensor given by binary
.
Note the binary
may have fewer elements than the
tensor size but, in such cases, it must strictly have
more elements than inspect_opts.limit