View Source Nx.Defn (Nx v0.2.0)

Numerical functions.

A numerical function is a subset of Elixir tailored for numerical computations. For example, the following function:

defmodule MyModule do
  import Nx.Defn

  defn softmax(t) do
    Nx.exp(t) / Nx.sum(Nx.exp(t))
  end
end

will work with scalars, vector, matrices, and n-dimensional tensors. Depending on your compiler of choice, the code can even be JIT-compiled and run either on the CPU or GPU.

To support these features, defn is a subset of Elixir. It replaces Elixir's Kernel by Nx.Defn.Kernel. Nx.Defn.Kernel provides tensor-aware operators, such as +, -, etc, while also preserving many high-level constructs known to Elixir developers, such as pipe operator, aliases, conditionals, pattern-matching, the access syntax, and more:

For example, the code above can also be written as:

defmodule MyModule do
  import Nx.Defn

  defn softmax(t) do
    t
    |> Nx.exp(t)
    |> then(& &1 / Nx.sum(&1))
  end
end

Please consult Nx.Defn.Kernel for a complete reference.

operators

Operators

defn attempts to keep as close to the Elixir semantics as possible but that's not achievable. For example, mathematical and bitwise operators (+, -, &&&, <<<, etc.) in Elixir work on numbers, which means mapping them to tensors is straight-forward and they largely preserve the same semantics, except they are now multi-dimensional.

On the other hand, the logical operators and, or, and not work with booleans in Elixir (true and false), which map to 0 and 1 in defn.

Therefore, when working with logical operators inside defn, 0 is considered false and all other numbers are considered true, which is represented as the number 1. For example, in defn, 0 and 1 as well as 0 and 2 return 0, while 1 and 1 or 1 and -1 will return 1.

The same semantics apply to conditional expressions inside defn, such as if, while, etc.

jit-compilers

JIT compilers

The power of Nx.Defn is given by its compilers. The default compiler is Nx.Defn.Evaluator, which executes the code in pure Elixir. You can use jit/3 to compile a function on the fly using a different compiler, such as EXLA:

Nx.Defn.jit(&MyModule.softmax/1, [my_tensor], compiler: EXLA)

The above will optimize, compile, and run softmax on the fly to the CPU (or the GPU) if available.

You can also change the default compiler for all numerical definitions (defn) by setting the default options. This can be done in your config/*.exs files as follows:

config :nx, :default_defn_options, compiler: EXLA

Now calling MyModule.softmax(my_tensor) will use EXLA even without wrapping it in jit/3. For scripts, you may also call Nx.Defn.global_default_options(compiler: EXLA).

defn functions are compiled when they are invoked, based on the type and shapes of the tensors given as arguments. The compilation is then cached based on the tensors shapes and types. Calling the same function with a tensor of different values but same shape and type means no recompilation is performed.

For those interested in writing custom compilers, see Nx.Defn.Compiler.

invoking-custom-elixir-code

Invoking custom Elixir code

Inside defn you can only call other defn functions and the functions in the Nx module. However, it is possible to use transforms to invoke any Elixir code:

defn add_and_mult(a, b, c) do
  res = a * b + c
  transform(res, &IO.inspect/1)
end

For example, the code above invokes &IO.inspect/1, which is not a defn function, with the value of res. This is useful as it allows developers to transform defn code to optimize, add new properties, and so on.

Transforms can also be used to manipulate Elixir data structures, such as options. defn expects all inputs to be tensors, with the exception of a default argument (declared with \\) which will be treated as options.

For example, imagine you want to support options where the :axis key is required. While you can't invoke Keyword directly, you can do it via a transform:

defn sum_axis(t, opts \\ []) do
  opts = keyword!(opts, [:axis])
  axis = transform(opts, &Keyword.fetch!(opts, :axis))
  Nx.sum(t, axes: [axis])
end

inputs-and-outputs-types

Inputs and outputs types

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

When numbers are given as arguments, they are always immediately converted to tensors on invocation. If you want to keep numbers as is or if you want to pass any other value to numerical definitions, they must be given as default arguments (see next subsection).

default-arguments

Default arguments

defn functions support default arguments. They are typically used as options. For example, imagine you want to create a function named zeros, which returns a tensor of zeroes with a given type and shape. It could be implemented like this:

defn zeros(opts \\ []) do
  opts = keyword!(opts, type: {:f, 32}, shape: {})
  Nx.broadcast(Nx.tensor(0, type: opts[:type]), opts[:shape])
end

The function above accepts opts which are then validated and given default values via the keyword!/2 function. Note that while it is possible to access options via the Access syntax, such as opts[:shape], it is not possible to directly call functions in the Keyword module inside defn. To freely manipulate any Elixir value inside defn, you have to use transforms, as described in the "Invoking custom Elixir code" section.

Additionally, defn supports anonymous as a direct input, without wrapping in a default argument.

Important! When it comes to JIT compilation, each different set of options and anonymous functions will lead to a different compilation of the numerical function.

Furthermore, if tensors are given through default arguments, they won't be cached effectively. Tensors in defn are cached based on their shape and type, not their value, but this is not true if the tensor is given via a default argument or captured by an anonymous function. For this reason, it is extremely discouraged to pass tensors through anonymous functions and default arguments.

working-with-maps-and-structs

Working with maps and structs

While Nx supports maps in defn, you must be careful if your numerical definitions are receiving maps and returning maps. For example, imagine this code:

defn update_a(map) do
  %{map | a: Nx.add(map.a, 1)}
end

The following code increments the value under the key :a by 1. However, because the function receives the whole map and returns the whole map, it means if the map has 120 keys, the whole map will be copied to the CPU/GPU, and then brought back.

However, if you do this instead:

defn update_a(map) do
  Nx.add(map.a, 1)
end

And then update the map on Elixir, outside of defn:

%{map | a: update_a(map)}

Nx will only send the parts of the map that matters.

Link to this section Summary

Functions

Gets the default options for the current process.

Sets the default options for defn in the current process.

Defines a public numerical function.

Defines a private numerical function.

Sets the default options globally.

Receives an anonymous function and returns a new anonymous function that returns the gradient of the input function when invoked.

Computes the gradient of the given var on fun.

Invokes the anonymous function with just-in-time compilation.

JITs the given function if outside of defn, otherwise invokes it.

Starts streaming the given anonymous function with just-in-time compilation.

Receives an anonymous function and returns a new anonymous function that returns the value and gradient of the input function when invoked.

Computes the value and gradient of the given var on fun with an optional data transformation.

Link to this section Functions

Gets the default options for the current process.

Link to this function

default_options(options)

View Source

Sets the default options for defn in the current process.

The options defined here apply to all future invocations of defn done by the current process. It also applies to calls to the jit/3 and stream/3 functions in this module.

The default options are stored only in the process dictionary and override any global options. This means if you start a separate process, such as Task, the default options must be set on the new process too.

This function is mostly used for scripting and testing. In your applications, you typically set the default options in your config files:

  config :nx, :default_defn_options, [compiler: EXLA, client: :cuda]
Link to this macro

defn(call, list)

View Source (macro)

Defines a public numerical function.

Link to this macro

defnp(call, list)

View Source (macro)

Defines a private numerical function.

Private numerical functions are always inlined by their callers at compilation time. This happens to all local function calls within defn.

Link to this function

global_default_options(options)

View Source

Sets the default options globally.

The options defined here apply to all future invocations of defn. It also applies to calls to the jit/3 and stream/3 functions in this module.

You must avoid calling this function at runtime. It is mostly useful during scripts or code notebooks to set a default. If you need to configure a global default options in your applications, you can do so in your config/*.exs files:

config :nx, :default_defn_options, [compiler: EXLA, client: :cuda]

Receives an anonymous function and returns a new anonymous function that returns the gradient of the input function when invoked.

examples

Examples

iex> fun = Nx.Defn.grad(fn x -> Nx.sin(x) end)
iex> fun.(Nx.tensor(0))
#Nx.Tensor<
  f32
  1.0
>

Computes the gradient of the given var on fun.

The result of the grad function must be a scalar tensor. If a non-scalar tensor is given, it is assumed the additional dimensions are batch dimensions.

examples

Examples

defn tanh_grad(t) do
  grad(t, &Nx.tanh/&1)
end

To differentiate on multiple vars, pass a tuple as first argument:

defn tanh_power_grad(a, b) do
  grad({a, b}, fn {a, b} -> Nx.tanh(a) + Nx.power(b, 2) end)
end

var_or_vars can be any Nx.Container with one or multiple tensors.

Link to this function

jit(fun, args, opts \\ [])

View Source

Invokes the anonymous function with just-in-time compilation.

The anonymous function will be invoked with tensor expressions which are JIT compiled and then invoked. For example, take the following definition:

defn softmax(t), do: Nx.exp(t) / Nx.sum(Nx.exp(t))

options

Options

Link to this function

jit_or_apply(fun, args, opts \\ [])

View Source

JITs the given function if outside of defn, otherwise invokes it.

It is not possible to invoke jit/3 inside defn, as all code inside defn is already jitted. However, some libraries may want to provide abstractions that can be invoked either inside defn or outside. In such cases, jit_or_apply/3 can be used to start jitting if it has been invoked outside of a numerical definition.

The opts are the same as the ones given to jit/3 and they are only used if invoking this function outside of defn.

Link to this function

stream(fun, args, opts \\ [])

View Source

Starts streaming the given anonymous function with just-in-time compilation.

At least two arguments are expected:

  1. The first argument is a tensor template of the data to be streamed in

  2. The second argument is a tensor with the stream initial state

The streaming function must return a two element tuple, the first element is the data to be sent and the second is the accumulator.

For each streamed chunk, you must call Nx.Stream.send/2 and Nx.Stream.recv/1. You don't need to call recv immediately after send, but doing so can be a useful mechanism to provide backpressure. Once all chunks are sent, you must use Nx.Stream.done/1 to receive the accumulated result. Let's see an example:

defmodule Streamed do
  import Nx.Defn

  defn sum(tensor, acc) do
    {acc, tensor + acc}
  end
end

Now let's invoke it:

stream = Nx.Defn.stream(&Streamed.sum/2, [Nx.template({}, {:s, 64}), 0])

for i <- 1..5 do
  Nx.Stream.send(stream, i)
  IO.inspect {:chunk, Nx.Stream.recv(stream)}
end

IO.inspect {:result, Nx.Stream.done(stream)}

It will print:

{:chunk, 0}
{:chunk, 1}
{:chunk, 2}
{:chunk, 3}
{:chunk, 4}
{:result, 5}

options

Options

Receives an anonymous function and returns a new anonymous function that returns the value and gradient of the input function when invoked.

examples

Examples

iex> fun = Nx.Defn.value_and_grad(fn x -> Nx.sin(x) end)
iex> {value, grad} = fun.(Nx.tensor(0))
iex> value
#Nx.Tensor<
  f32
  0.0
>
iex> grad
#Nx.Tensor<
  f32
  1.0
>
Link to this function

value_and_grad(var_or_vars, fun, transform \\ & &1)

View Source

Computes the value and gradient of the given var on fun with an optional data transformation.

It returns a tuple with the value and the gradient.

examples

Examples

defn tanh_grad(t) do
  value_and_grad(t, &Nx.tanh/&1)
end

To differentiate on multiple vars, pass a tuple as first argument:

defn tanh_power_grad(a, b) do
  value_and_grad({a, b}, fn {a, b} -> Nx.tanh(a) + Nx.power(b, 2) end)
end

var_or_vars can be any Nx.Container with one or multiple tensors.

transform allows you to transform the expression before the gradient is calculated. This enables optimizations that reuse parts of expressions. As an example, consider the following objective function:

defn objective(predict_fn, loss_fn, params, inputs, targets) do
  preds = predict_fn.(params, inputs)
  loss = loss_fn.(preds, targets)
  {preds, loss}
end

You can compute the gradient with respect to just the loss function by applying a transform:

{{preds, loss}, gradient} = value_and_grad(params, &objective(predict_fn, loss_fn, &1, inputs, targets), &elem(&1, 1))

preds can be re-used to compute other metrics such as accuracy, absolute error, etc. without having to do another forward pass.