defmodule Nx do

@moduledoc """

Numerical Elixir.

The `Nx` library is a collection of functions and data

types to work with Numerical Elixir. This module defines

the main entry point for building and working with said

data-structures. For example, to create an n-dimensional

tensor, do:

iex> t = Nx.tensor([[1, 2], [3, 4]])

iex> Nx.shape(t)

{2, 2}

`Nx` also provides the so-called numerical definitions under

the `Nx.Defn` module. They are a subset of Elixir tailored for

numerical computations. For example, it overrides Elixir's

default operators so they are tensor-aware:

defn softmax(t) do

Nx.exp(t) / Nx.sum(Nx.exp(t))

end

Code inside `defn` functions can also be given to custom compilers,

which can compile said functions just-in-time (JIT) to run on the

CPU or on the GPU.

## References

Here is a general outline of the main references in this library:

* For an introduction, see our [Intro to Nx](intro-to-nx.livemd) guide

* This module provides the main API for working with tensors

* `Nx.Defn` provides numerical definitions, CPU/GPU compilation, gradients, and more

* `Nx.LinAlg` provides functions related to linear algebra

* `Nx.Constants` declares many constants commonly used in numerical code

Continue reading this documentation for an overview of creating,

broadcasting, and accessing/slicing Nx tensors.

## Creating tensors

The main APIs for creating tensors are `tensor/2`, `from_binary/2`,

`iota/2`, `eye/2`, and `broadcast/3`.

The tensor types can be one of:

* unsigned integers (`u8`, `u16`, `u32`, `u64`)

* signed integers (`s8`, `s16`, `s32`, `s64`)

* floats (`f16`, `f32`, `f64`)

* brain floats (`bf16`)

* and complex numbers (`c64`, `c128`)

The types are tracked as tuples:

iex> Nx.tensor([1, 2, 3], type: {:f, 32})

#Nx.Tensor<

f32[3]

[1.0, 2.0, 3.0]

>

But a shortcut atom notation is also available:

iex> Nx.tensor([1, 2, 3], type: :f32)

#Nx.Tensor<

f32[3]

[1.0, 2.0, 3.0]

>

The tensor dimensions can also be named, via the `:names` option

available to all creation functions:

iex> Nx.iota({2, 3}, names: [:x, :y])

#Nx.Tensor<

s64[x: 2][y: 3]

[

[0, 1, 2],

[3, 4, 5]

]

>

Finally, for creating vectors and matrices, a sigil notation

is available:

iex> import Nx, only: :sigils

iex> ~V[1 2 3]f32

#Nx.Tensor<

f32[3]

[1.0, 2.0, 3.0]

>

iex> import Nx, only: :sigils

iex> ~M'''

...> 1 2 3

...> 4 5 6

...> '''s32

#Nx.Tensor<

s32[2][3]

[

[1, 2, 3],

[4, 5, 6]

]

>

All other APIs accept exclusively numbers or tensors, unless

explicitly noted otherwise.

## Broadcasting

Broadcasting allows operations on two tensors of different shapes

to match. For example, most often operations between tensors have

the same shape:

iex> a = Nx.tensor([1, 2, 3])

iex> b = Nx.tensor([10, 20, 30])

iex> Nx.add(a, b)

#Nx.Tensor<

s64[3]

[11, 22, 33]

>

Now let's imagine you want to multiply a large tensor of dimensions

1000x1000x1000 by 2. If you had to create a similarly large tensor

only to perform this operation, it would be inefficient. Therefore,

you can simply multiply this large tensor by the scalar 2, and Nx

will propagate its dimensions at the time the operation happens,

without allocating a large intermediate tensor:

iex> Nx.multiply(Nx.tensor([1, 2, 3]), 2)

#Nx.Tensor<

s64[3]

[2, 4, 6]

>

In practice, broadcasting is not restricted only to scalars; it

is a general algorithm that applies to all dimensions of a tensor.

When broadcasting, `Nx` compares the shapes of the two tensors,

starting with the trailing ones, such that:

* If the dimensions have equal size, then they are compatible

* If one of the dimensions have size of 1, it is "broadcast"

to match the dimension of the other

In case one tensor has more dimensions than the other, the missing

dimensions are considered to be of size one. Here are some examples

of how broadcast would work when multiplying two tensors with the

following shapes:

s64[3] * s64

#=> s64[3]

s64[255][255][3] * s64[3]

#=> s64[255][255][3]

s64[2][1] * s[1][2]

#=> s64[2][2]

s64[5][1][4][1] * s64[3][4][5]

#=> s64[5][3][4][5]

If any of the dimensions do not match or are not 1, an error is

raised.

## Access syntax (slicing)

Nx tensors implement Elixir's access syntax. This allows developers

to slice tensors up and easily access sub-dimensions and values.

Access accepts integers:

iex> t = Nx.tensor([[1, 2], [3, 4]])

iex> t[0]

#Nx.Tensor<

s64[2]

[1, 2]

>

iex> t[1]

#Nx.Tensor<

s64[2]

[3, 4]

>

iex> t[1][1]

#Nx.Tensor<

s64

4

>

If a negative index is given, it accesses the element from the back:

iex> t = Nx.tensor([[1, 2], [3, 4]])

iex> t[-1][-1]

#Nx.Tensor<

s64

4

>

Out of bound access will raise:

iex> Nx.tensor([1, 2])[2]

** (ArgumentError) index 2 is out of bounds for axis 0 in shape {2}

iex> Nx.tensor([1, 2])[-3]

** (ArgumentError) index -3 is out of bounds for axis 0 in shape {2}

The index can also be another tensor. If the tensor is a scalar, it must

be a value between 0 and the dimension size, and it behaves the same as

an integer. Out of bound dynamic indexes are always clamped to the tensor

dimensions:

iex> two = Nx.tensor(2)

iex> t = Nx.tensor([[1, 2], [3, 4]])

iex> t[two][two]

#Nx.Tensor<

s64

4

>

For example, a `minus_one` dynamic index will be clamped to zero:

iex> minus_one = Nx.tensor(-1)

iex> t = Nx.tensor([[1, 2], [3, 4]])

iex> t[minus_one][minus_one]

#Nx.Tensor<

s64

1

>

A multi-dimensional tensor uses its values to fetch the leading

dimension of the tensor, placing them within the shape of the

indexing tensor. It is equivalent to `take/3`:

iex> t = Nx.tensor([[1, 2], [3, 4]])

iex> t[Nx.tensor([1, 0])]

#Nx.Tensor<

s64[2][2]

[

[3, 4],

[1, 2]

]

>

The example shows how the retrieved indexes are nested

with the accessed shape and that you may also access

repeated indices:

iex> t = Nx.tensor([[1, 2], [3, 4]])

iex> t[Nx.tensor([[1, 0, 1]])]

#Nx.Tensor<

s64[1][3][2]

[

[

[3, 4],

[1, 2],

[3, 4]

]

]

>

Access also accepts ranges. Ranges in Elixir are inclusive:

iex> t = Nx.tensor([[1, 2], [3, 4], [5, 6], [7, 8]])

iex> t[0..1]

#Nx.Tensor<

s64[2][2]

[

[1, 2],

[3, 4]

]

>

Ranges can receive negative positions and they will read from

the back. In such cases, the range step must be explicitly given

and the right-side of the range must be equal or greater than

the left-side:

iex> t = Nx.tensor([[1, 2], [3, 4], [5, 6], [7, 8]])

iex> t[1..-2//1]

#Nx.Tensor<

s64[2][2]

[

[3, 4],

[5, 6]

]

>

As you can see, accessing with a range does not eliminate the

accessed axis. This means that, if you try to cascade ranges,

you will always be filtering the highest dimension:

iex> t = Nx.tensor([[1, 2], [3, 4], [5, 6], [7, 8]])

iex> t[1..-1//1] # Drop the first "row"

#Nx.Tensor<

s64[3][2]

[

[3, 4],

[5, 6],

[7, 8]

]

>

iex> t[1..-1//1][1..-1//1] # Drop the first "row" twice

#Nx.Tensor<

s64[2][2]

[

[5, 6],

[7, 8]

]

>

Therefore, if you want to slice across multiple dimensions, you can wrap

the ranges in a list:

iex> t = Nx.tensor([[1, 2], [3, 4], [5, 6], [7, 8]])

iex> t[[1..-1//1, 1..-1//1]] # Drop the first "row" and the first "column"

#Nx.Tensor<

s64[3][1]

[

[4],

[6],

[8]

]

>

You can also use `..` as the full-slice range, which means you want to

keep a given dimension as is:

iex> t = Nx.tensor([[1, 2], [3, 4], [5, 6], [7, 8]])

iex> t[[.., 1..-1//1]] # Drop only the first "column"

#Nx.Tensor<

s64[4][1]

[

[2],

[4],

[6],

[8]

]

>

You can mix both ranges and integers in the list too:

iex> t = Nx.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]])

iex> t[[1..2, 2]]

#Nx.Tensor<

s64[2]

[6, 9]

>

If the list has less elements than axes, the remaining dimensions

are returned in full:

iex> t = Nx.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]])

iex> t[[1..2]]

#Nx.Tensor<

s64[2][3]

[

[4, 5, 6],

[7, 8, 9]

]

>

The access syntax also pairs nicely with named tensors. By using named

tensors, you can pass only the axis you want to slice, leaving the other

axes intact:

iex> t = Nx.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]], names: [:x, :y])

iex> t[x: 1..2]

#Nx.Tensor<

s64[x: 2][y: 3]

[

[4, 5, 6],

[7, 8, 9]

]

>

iex> t[x: 1..2, y: 0..1]

#Nx.Tensor<

s64[x: 2][y: 2]

[

[4, 5],

[7, 8]

]

>

iex> t[x: 1, y: 0..1]

#Nx.Tensor<

s64[y: 2]

[4, 5]

>

For a more complex slicing rules, including strides, you

can always fallback to `Nx.slice/4`.

## Backends

The `Nx` library has built-in support for multiple backends.

A tensor is always handled by a backend, the default backend

being `Nx.BinaryBackend`, which means the tensor is allocated

as a binary within the Erlang VM.

Most often backends are used to provide a completely different

implementation of tensor operations, often accelerated to the GPU.

In such cases, you want to guarantee all tensors are allocated in

the new backend. This can be done by configuring your runtime:

# config/runtime.exs

import Config

config :nx, default_backend: EXLA.Backend

In your notebooks and on `Mix.install/2`, you might:

Mix.install(

[

{:nx, ">= 0.0.0"}

],

config: [nx: [default_backend: EXLA.Backend]]

)

Or by calling `Nx.global_default_backend/1` (less preferrable):

Nx.global_default_backend(EXLA.Backend)

To pass options to the backend, replacing `EXLA.Backend` by

`{EXLA.Backend, client: :cuda}` or similar. See the documentation

for [EXLA](https://hexdocs.pm/exla) and [Torchx](https://hexdocs.pm/torchx)

for installation and GPU support.

To implement your own backend, check the `Nx.Tensor` behaviour.

"""

import Nx.Shared

import Nx.Defn.Kernel, only: [keyword!: 2]

alias Nx.Tensor, as: T

@typedoc """

Represents a numerical value.

Can be a plain number, a `Complex` number or an `Nx.Tensor`.

See also: `is_tensor/1`

"""

@type t :: number | Complex.t() | Nx.Tensor.t()

@type shape :: number() | Nx.Tensor.t() | Nx.Tensor.shape()

@type axis :: Nx.Tensor.axis()

@type axes :: Nx.Tensor.axes()

@type template :: Nx.Tensor.t(%Nx.TemplateBackend{})

@file_prefix <<?n, ?x>>

@file_version 1

@non_finite [:neg_infinity, :infinity, :nan]

@doc """

Checks whether the value is a valid numerical value.

Returns true if the value is a `number`, a non-finite atom (like `:infinity`),

a `Complex` number or an `Nx.Tensor`.

See also: `t:t/0`

"""

@doc type: :guards

defguard is_tensor(t)

when is_number(t) or is_struct(t, T) or is_struct(t, Complex) or t in @non_finite

## Creation API

@doc """

Builds a tensor.

The argument must be one of:

* a tensor

* a number (which means the tensor is scalar/zero-dimensional)

* a boolean (also scalar/zero-dimensional)

* an arbitrarily nested list of numbers and booleans

If a new tensor has to be allocated, it will be allocated in

`Nx.default_backend/0`, unless the `:backend` option is given,

which overrides the default one.

## Examples

A number returns a tensor of zero dimensions:

iex> Nx.tensor(0)

#Nx.Tensor<

s64

0

>

iex> Nx.tensor(1.0)

#Nx.Tensor<

f32

1.0

>

Giving a list returns a vector (a one-dimensional tensor):

iex> Nx.tensor([1, 2, 3])

#Nx.Tensor<

s64[3]

[1, 2, 3]

>

iex> Nx.tensor([1.2, 2.3, 3.4, 4.5])

#Nx.Tensor<

f32[4]

[1.2000000476837158, 2.299999952316284, 3.4000000953674316, 4.5]

>

The type can be explicitly given. Integers and floats

bigger than the given size overflow:

iex> Nx.tensor([300, 301, 302], type: :s8)

#Nx.Tensor<

s8[3]

[44, 45, 46]

>

Mixed types give higher priority to floats:

iex> Nx.tensor([1, 2, 3.0])

#Nx.Tensor<

f32[3]

[1.0, 2.0, 3.0]

>

Boolean values are also accepted, where `true` is

converted to `1` and `false` to `0`, with the type

being inferred as `{:u, 8}`

iex> Nx.tensor(true)

#Nx.Tensor<

u8

1

>

iex> Nx.tensor(false)

#Nx.Tensor<

u8

0

>

iex> Nx.tensor([true, false])

#Nx.Tensor<

u8[2]

[1, 0]

>

Multi-dimensional tensors are also possible:

iex> Nx.tensor([[1, 2, 3], [4, 5, 6]])

#Nx.Tensor<

s64[2][3]

[

[1, 2, 3],

[4, 5, 6]

]

>

iex> Nx.tensor([[1, 2], [3, 4], [5, 6]])

#Nx.Tensor<

s64[3][2]

[

[1, 2],

[3, 4],

[5, 6]

]

>

iex> Nx.tensor([[[1, 2], [3, 4], [5, 6]], [[-1, -2], [-3, -4], [-5, -6]]])

#Nx.Tensor<

s64[2][3][2]

[

[

[1, 2],

[3, 4],

[5, 6]

],

[

[-1, -2],

[-3, -4],

[-5, -6]

]

]

>

## Floats and complex numbers

Besides single-precision (32 bits), floats can also have

half-precision (16) or double-precision (64):

iex> Nx.tensor([1, 2, 3], type: :f16)

#Nx.Tensor<

f16[3]

[1.0, 2.0, 3.0]

>

iex> Nx.tensor([1, 2, 3], type: :f64)

#Nx.Tensor<

f64[3]

[1.0, 2.0, 3.0]

>

Brain-floating points are also supported:

iex> Nx.tensor([1, 2, 3], type: :bf16)

#Nx.Tensor<

bf16[3]

[1.0, 2.0, 3.0]

>

In all cases, the non-finite values negative infinity (-Inf),

infinity (Inf), and "not a number" (NaN) can be represented by

the atoms `:neg_infinity`, `:infinity`, and `:nan` respectively:

iex> Nx.tensor([:neg_infinity, :nan, :infinity])

#Nx.Tensor<

f32[3]

[-Inf, NaN, Inf]

>

Finally, complex numbers are also supported in tensors:

iex> Nx.tensor(Complex.new(1, -1))

#Nx.Tensor<

c64

1.0-1.0i

>

## Naming dimensions

You can provide names for tensor dimensions. Names are atoms:

iex> Nx.tensor([[1, 2, 3], [4, 5, 6]], names: [:x, :y])

#Nx.Tensor<

s64[x: 2][y: 3]

[

[1, 2, 3],

[4, 5, 6]

]

>

Names make your code more expressive:

iex> Nx.tensor([[[1, 2, 3], [4, 5, 6], [7, 8, 9]]], names: [:batch, :height, :width])

#Nx.Tensor<

s64[batch: 1][height: 3][width: 3]

[

[

[1, 2, 3],

[4, 5, 6],

[7, 8, 9]

]

]

>

You can also leave dimension names as `nil`:

iex> Nx.tensor([[[1, 2, 3], [4, 5, 6], [7, 8, 9]]], names: [:batch, nil, nil])

#Nx.Tensor<

s64[batch: 1][3][3]

[

[

[1, 2, 3],

[4, 5, 6],

[7, 8, 9]

]

]

>

However, you must provide a name for every dimension in the tensor:

iex> Nx.tensor([[[1, 2, 3], [4, 5, 6], [7, 8, 9]]], names: [:batch])

** (ArgumentError) invalid names for tensor of rank 3, when specifying names every dimension must have a name or be nil

## Tensors

Tensors can also be given as inputs:

iex> Nx.tensor(Nx.tensor([1, 2, 3]))

#Nx.Tensor<

s64[3]

[1, 2, 3]

>

If the `:backend` and `:type` options are given, the tensor will

compared against those values and raise in case of mismatch:

iex> Nx.tensor(Nx.tensor([1, 2, 3]), type: :f32)

** (ArgumentError) Nx.tensor/2 expects a tensor with type :f32 but it was given a tensor of type {:s, 64}

The `:backend` option will check only against the backend name

and not specific backend configuration such as device and client.

In case the backend differs, it will also raise.

The names in the given tensor are always discarded but Nx will raise

in case the tensor already has names that conflict with the assigned ones:

iex> Nx.tensor(Nx.tensor([1, 2, 3]), names: [:row])

#Nx.Tensor<

s64[row: 3]

[1, 2, 3]

>

iex> Nx.tensor(Nx.tensor([1, 2, 3], names: [:column]))

#Nx.Tensor<

s64[3]

[1, 2, 3]

>

iex> Nx.tensor(Nx.tensor([1, 2, 3], names: [:column]), names: [:row])

** (ArgumentError) cannot merge name :column on axis 0 with name :row on axis 0

## Options

* `:type` - sets the type of the tensor. If one is not given,

one is automatically inferred based on the input.

* `:names` - dimension names. If you wish to specify dimension

names you must specify a name for every dimension in the tensor.

Only `nil` and atoms are supported as dimension names.

* `:backend` - the backend to allocate the tensor on. It is either

an atom or a tuple in the shape `{backend, options}`. It defaults

to `Nx.default_backend/0` for new tensors

"""

@doc type: :creation

def tensor(arg, opts \\ [])

def tensor(%Nx.Tensor{} = tensor, opts) do

opts = keyword!(opts, [:type, :names, :backend])

tensor =

if backend = opts[:backend] do

case backend!(backend) do

{backend, _options} when tensor.data.__struct__ == backend ->

tensor

{backend, _} ->

raise ArgumentError,

"Nx.tensor/2 wants to allocate on backend #{inspect(backend)} " <>

"but it was given a tensor allocated on #{inspect(tensor.data.__struct__)}"

end

else

tensor

end

tensor =

if type = opts[:type] do

if tensor.type == Nx.Type.normalize!(type) do

tensor

else

raise ArgumentError,

"Nx.tensor/2 expects a tensor with type #{inspect(type)} " <>

"but it was given a tensor of type #{inspect(tensor.type)}"

end

else

tensor

end

# We merge to check for conflicts but ultimately discard the tensor.names for consistency

names =

if names = opts[:names] do

names = Nx.Shape.named_axes!(names, tensor.shape)

_ = Nx.Shape.merge_names!(tensor.names, names)

names

else

List.duplicate(nil, tuple_size(tensor.shape))

end

%{tensor | names: names}

end

def tensor(arg, opts) do

opts = keyword!(opts, [:type, :names, :backend])

type = Nx.Type.normalize!(opts[:type] || infer_type(arg))

tensor(arg, type, opts)

end

defp infer_type([head | tail]) when is_list(tail) do

Enum.reduce(tail, infer_type(head), &Nx.Type.merge(infer_type(&1), &2))

end

defp infer_type(number)

when is_number(number) or is_struct(number, Complex) or number in @non_finite or

is_boolean(number) do

Nx.Type.infer(number)

end

defp infer_type(%Nx.Tensor{} = value) do

raise ArgumentError,

"invalid value given to Nx.tensor/1. If you want to create a tensor from other tensors, " <>

"consider using Nx.concatenate/2 or Nx.stack/2 instead. Got: #{inspect(value)}"

end

defp infer_type(value) do

raise ArgumentError, "invalid value given to Nx.tensor/1, got: #{inspect(value)}"

end

defp tensor(true, type, opts), do: tensor(1, type, opts)

defp tensor(false, type, opts), do: tensor(0, type, opts)

defp tensor(arg, type, opts) when is_number(arg) do

names = Nx.Shape.named_axes!(opts[:names], {})

{backend, backend_options} = backend_from_options!(opts) || default_backend()

backend.constant(%T{shape: {}, type: type, names: names}, arg, backend_options)

end

defp tensor(%Complex{} = arg, {:c, size}, opts) do

names = Nx.Shape.named_axes!(opts[:names], {})

{backend, backend_options} = backend_from_options!(opts) || default_backend()

backend.constant(%T{shape: {}, type: {:c, size}, names: names}, arg, backend_options)

end

defp tensor(%Complex{}, type, _) do

raise ArgumentError,

"invalid type for complex number. Expected {:c, 64} or {:c, 128}, got: #{inspect(type)}"

end

defp tensor(arg, type, opts) when arg in @non_finite do

names = Nx.Shape.named_axes!(opts[:names], {})

{backend, backend_options} = backend_from_options!(opts) || default_backend()

data = number_to_binary(arg, type)

backend.from_binary(%T{shape: {}, type: type, names: names}, data, backend_options)

end

defp tensor(arg, type, opts) when is_list(arg) do

{shape, data} = flatten_list(arg, type)

if data == "" do

raise "cannot build empty tensor"

end

names = Nx.Shape.named_axes!(opts[:names], shape)

{backend, backend_options} = backend_from_options!(opts) || default_backend()

backend.from_binary(%T{shape: shape, type: type, names: names}, data, backend_options)

end

defp flatten_list(list, type) do

{dimensions, acc} = flatten_list(list, type, [], [])

{dimensions |> Enum.reverse() |> List.to_tuple(),

acc |> Enum.reverse() |> :erlang.list_to_binary()}

end

defp flatten_list([], _type, dimensions, acc) do

{[0 | dimensions], acc}

end

defp flatten_list([head | rest], type, parent_dimensions, acc) when is_list(head) do

{child_dimensions, acc} = flatten_list(head, type, [], acc)

{n, acc} =

Enum.reduce(rest, {1, acc}, fn list, {count, acc} ->

case flatten_list(list, type, [], acc) do

{^child_dimensions, acc} ->

{count + 1, acc}

{other_dimensions, _acc} ->

raise ArgumentError,

"cannot build tensor because lists have different shapes, got " <>

inspect(List.to_tuple(child_dimensions)) <>

" at position 0 and " <>

inspect(List.to_tuple(other_dimensions)) <> " at position #{count + 1}"

end

end)

{child_dimensions ++ [n | parent_dimensions], acc}

end

defp flatten_list(list, type, dimensions, acc) do

{[length(list) | dimensions],

Enum.reduce(list, acc, &[tensor_or_number_to_binary(&1, type) | &2])}

end

defp tensor_or_number_to_binary(true, type), do: tensor_or_number_to_binary(1, type)

defp tensor_or_number_to_binary(false, type), do: tensor_or_number_to_binary(0, type)

defp tensor_or_number_to_binary(number, type)

when is_number(number)

when is_struct(number, Complex)

when number in @non_finite do

number_to_binary(number, type)

end

defp tensor_or_number_to_binary(value, _type) do

raise ArgumentError, "invalid value given to Nx.tensor/1, got: #{inspect(value)}"

end

@doc """

Creates a tensor template.

You can't perform any operation on this tensor.

It exists exclusively to define APIs that say

a tensor with a certain type, shape, and names

is expected in the future.

## Examples

iex> Nx.template({2, 3}, :f32)

#Nx.Tensor<

f32[2][3]

Nx.TemplateBackend

>

iex> Nx.template({2, 3}, {:f, 32}, names: [:rows, :columns])

#Nx.Tensor<

f32[rows: 2][columns: 3]

Nx.TemplateBackend

>

Although note it is impossible to perform any operation on a tensor template:

iex> t = Nx.template({2, 3}, {:f, 32}, names: [:rows, :columns])

iex> Nx.abs(t)

** (RuntimeError) cannot perform operations on a Nx.TemplateBackend tensor

To convert existing tensors to templates, use `to_template/1`.

"""

@doc type: :creation

def template(shape, type, opts \\ []) when is_tuple(shape) do

opts = keyword!(opts, [:names])

type = Nx.Type.normalize!(type)

names = Nx.Shape.named_axes!(opts[:names], shape)

%T{shape: shape, type: type, names: names, data: %Nx.TemplateBackend{}}

end

for t <- [:u8, :u16, :u32, :u64, :s8, :s16, :s32, :s64, :bf16, :f16, :f32, :f64] do

@doc """

Short-hand function for creating tensor of type `#{t}`.

This is just an alias for `Nx.tensor(tensor, type: #{t})`.

"""

@doc type: :creation

def unquote(t)(tensor), do: Nx.tensor(tensor, type: unquote(t))

end

@doc """

Converts a tensor (or tuples and maps of tensors) to tensor templates.

Templates are useful when you need to pass types and shapes to

operations and the data is not yet available.

For convenience, this function accepts tensors and any container

(such as maps and tuples as defined by the `Nx.LazyContainer` protocol)

and recursively converts all tensors to templates.

## Examples

iex> Nx.iota({2, 3}) |> Nx.to_template()

#Nx.Tensor<

s64[2][3]

Nx.TemplateBackend

>

iex> {int, float} = Nx.to_template({1, 2.0})

iex> int

#Nx.Tensor<

s64

Nx.TemplateBackend

>

iex> float

#Nx.Tensor<

f32

Nx.TemplateBackend

>

Although note it is impossible to perform any operation on a tensor template:

iex> t = Nx.iota({2, 3}) |> Nx.to_template()

iex> Nx.abs(t)

** (RuntimeError) cannot perform operations on a Nx.TemplateBackend tensor

To build a template from scratch, use `template/3`.

"""

@doc type: :conversion

def to_template(tensor_or_container) do

tensor_or_container

|> Nx.LazyContainer.traverse(:ok, fn template, _fun, :ok -> {template, :ok} end)

|> then(fn {template, :ok} -> template end)

end

@doc """

Creates a tensor with the given shape which increments

along the provided axis. You may optionally provide dimension

names.

If no axis is provided, index counts up at each element.

If a tensor or a number are given, the shape and names are taken from the tensor.

## Options

* `:type` - the type of the tensor

* `:axis` - an axis to repeat the iota over