nx

defmodule Nx.Backend do
  @moduledoc """
  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 the `Nx` 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 by `defn` to build
      expression trees that are traversed by custom compilers.

  This module also includes functions that are meant to be shared
  across backends.

  Backends may publish backend documentation guides that mirror the Nx API and
  describe divergent behaviour, backend-specific options, and limitations — see
  the [backend documentation convention](backend_documentation-convention.html)
  guide.
  """

  @type t :: %{__struct__: atom()}

  @type tensor :: Nx.Tensor.t()
  @type shape :: Nx.Tensor.shape()
  @type axis :: Nx.Tensor.axis()
  @type axes :: Nx.Tensor.axes()
  @type backend_options :: term()

  @callback init(keyword()) :: backend_options

  @callback constant(out :: tensor, number | Complex.t(), backend_options) :: tensor
  @callback from_binary(out :: tensor, binary, backend_options) :: tensor
  @callback eye(tensor, backend_options) :: tensor
  @callback iota(tensor, axis | nil, backend_options) :: tensor

  @callback backend_deallocate(tensor) :: :ok | :already_deallocated
  @callback backend_copy(tensor, module, backend_options) :: tensor
  @callback backend_transfer(tensor, module, backend_options) :: tensor
  @callback to_batched(out :: tensor, tensor, keyword) :: [tensor]
  @callback to_binary(tensor, limit :: non_neg_integer) :: binary
  @callback inspect(tensor, Inspect.Opts.t()) :: tensor
  @callback from_pointer(
              opaque_pointer :: term(),
              type :: tuple(),
              shape :: tuple(),
              backend_opts :: keyword(),
              opts :: keyword()
            ) :: tensor | no_return()
  @callback to_pointer(tensor, opts :: keyword) :: term() | no_return()

  @callback as_type(out :: tensor, tensor) :: tensor
  @callback bitcast(out :: tensor, tensor) :: tensor
  @callback reshape(out :: tensor, tensor) :: tensor
  @callback squeeze(out :: tensor, tensor, axes) :: tensor
  @callback broadcast(out :: tensor, tensor, shape, axes) :: tensor
  @callback transpose(out :: tensor, tensor, axes) :: tensor
  @callback pad(out :: tensor, tensor, pad_value :: tensor, padding_config :: list()) :: tensor
  @callback reverse(out :: tensor, tensor, axes) :: tensor

  @callback dot(out :: tensor, tensor, axes, axes, tensor, axes, axes) :: tensor
  @callback clip(out :: tensor, tensor, min :: tensor, max :: tensor) :: tensor
  @callback slice(out :: tensor, tensor, list, list, list) :: tensor
  @callback put_slice(out :: tensor, tensor, tensor, list) :: tensor
  @callback gather(out :: tensor, input :: tensor, indices :: tensor, keyword) :: tensor
  @callback concatenate(out :: tensor, tensor, axis) :: tensor
  @callback stack(out :: tensor, tensor, axis) :: tensor
  @callback select(out :: tensor, tensor, tensor, tensor) :: tensor

  @callback conv(out :: tensor, tensor, kernel :: tensor, keyword) :: tensor
  @callback all(out :: tensor, tensor, keyword) :: tensor
  @callback any(out :: tensor, tensor, keyword) :: tensor
  @callback sum(out :: tensor, tensor, keyword) :: tensor
  @callback product(out :: tensor, tensor, keyword) :: tensor
  @callback reduce_max(out :: tensor, tensor, keyword) :: tensor
  @callback reduce_min(out :: tensor, tensor, keyword) :: tensor
  @callback argmax(out :: tensor, tensor, keyword) :: tensor
  @callback argmin(out :: tensor, tensor, keyword) :: tensor
  @callback reduce(out :: tensor, tensor, acc :: tensor, keyword, fun) :: tensor
  @callback window_reduce(out :: tensor, tensor, acc :: tensor, shape, keyword, fun) :: tensor
  @callback window_sum(out :: tensor, tensor, shape, keyword) :: tensor
  @callback window_product(out :: tensor, tensor, shape, keyword) :: tensor
  @callback window_max(out :: tensor, tensor, shape, keyword) :: tensor
  @callback window_min(out :: tensor, tensor, shape, keyword) :: tensor
  @callback sort(out :: tensor, tensor, keyword) :: tensor
  @callback argsort(out :: tensor, tensor, keyword) :: tensor
  @callback window_scatter_max(out :: tensor, tensor, tensor, tensor, shape, keyword) :: tensor
  @callback window_scatter_min(out :: tensor, tensor, tensor, tensor, shape, keyword) :: tensor
  @callback indexed_add(out :: tensor, tensor, indices :: tensor, updates :: tensor, keyword) ::
              tensor
  @callback indexed_put(out :: tensor, tensor, indices :: tensor, updates :: tensor, keyword) ::
              tensor

  @callback triangular_solve(out :: tensor, a :: tensor, b :: tensor, keyword) :: tensor

  @callback fft(out :: tensor, tensor, keyword) :: tensor
  @callback ifft(out :: tensor, tensor, keyword) :: tensor

  binary_ops =
    [:add, :subtract, :multiply, :pow, :remainder, :divide, :atan2, :min, :max, :quotient] ++
      [:bitwise_and, :bitwise_or, :bitwise_xor, :left_shift, :right_shift] ++
      [:equal, :not_equal, :greater, :less, :greater_equal, :less_equal] ++
      [:logical_and, :logical_or, :logical_xor]

  for binary_op <- binary_ops do
    @callback unquote(binary_op)(out :: tensor, tensor, tensor) :: tensor
  end

  unary_ops =
    Enum.map(Nx.Shared.unary_math_funs(), &elem(&1, 0)) ++
      [:abs, :bitwise_not, :ceil, :conjugate, :floor, :negate, :round, :sign] ++
      [:count_leading_zeros, :population_count, :real, :imag, :is_nan, :is_infinity]

  for unary_op <- unary_ops do
    @callback unquote(unary_op)(out :: tensor, tensor) :: tensor
  end

  ## Block extension

  @doc """
  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 block(struct, output :: tensor | tuple, args :: [term], fun) :: tensor | tuple

  ## Inspect implementation

  require Nx.Shared
  alias Inspect.Algebra, as: IA

  @doc """
  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`.
  """
  def inspect(%{shape: shape, type: type}, binary, inspect_opts) do
    open = IA.color("[", :list, inspect_opts)
    sep = IA.color(",", :list, inspect_opts)
    close = IA.color("]", :list, inspect_opts)

    dims = Tuple.to_list(shape)

    {data, _rest, _limit} =
      chunk(dims, binary, type, inspect_opts.limit, {open, sep, close})

    data
  end

  defp chunk([], data, type, limit, _docs) do
    {doc, tail} =
      case type do
        {:s, size} ->
          <<seg::size(^size)-signed-integer-native, tail::binary>> = data
          {Integer.to_string(seg), tail}

        {:u, size} ->
          <<seg::size(^size)-unsigned-integer-native, tail::binary>> = data
          {Integer.to_string(seg), tail}

        {:c, size} ->
          half = div(size, 2)
          <<re::size(^half)-bitstring, im::size(^half)-bitstring, tail::binary>> = data
          {complex_to_string(re, im, half), tail}

        {type, size} ->
          <<float::size(^size)-bitstring, tail::bitstring>> = data
          {float_to_string(float, type, size), tail}
      end

    if limit == :infinity, do: {doc, tail, limit}, else: {doc, tail, limit - 1}
  end

  defp chunk([dim | dims], data, type, limit, {open, sep, close} = docs) do
    {acc, rest, limit} =
      chunk_each(dim, data, [], limit, fn chunk, limit ->
        chunk(dims, chunk, type, limit, docs)
      end)

    {open, sep, close, nest} =
      if dims == [] do
        {open, IA.concat(sep, " "), close, 0}
      else
        {IA.concat(open, IA.line()), IA.concat(sep, IA.line()), IA.concat(IA.line(), close), 2}
      end

    doc =
      open
      |> IA.concat(IA.concat(Enum.intersperse(acc, sep)))
      |> IA.nest(nest)
      |> IA.concat(close)

    {doc, rest, limit}
  end

  defp chunk_each(0, data, acc, limit, _fun) do
    {Enum.reverse(acc), data, limit}
  end

  defp chunk_each(_dim, data, acc, 0, _fun) do
    {Enum.reverse(["..." | acc]), data, 0}
  end

  defp chunk_each(dim, data, acc, limit, fun) do
    {doc, rest, limit} = fun.(data, limit)
    chunk_each(dim - 1, rest, [doc | acc], limit, fun)
  end

  defp float_to_string(<<float::64-float-native>>, :f, 64),
    do: Float.to_string(float)

  defp float_to_string(binary, :f, 64),
    do:
      (case Nx.Shared.read_non_finite(binary, 64) do
         :nan -> "NaN"
         :infinity -> "Inf"
         :neg_infinity -> "-Inf"
       end)

  defp float_to_string(<<bits::32-native-unsigned>>, :f, 32),
    do: Nx.Ryu.bits_to_decimal(bits, 23, 8)

  defp float_to_string(<<bits::16-native-unsigned>>, :f, 16),
    do: Nx.Ryu.bits_to_decimal(bits, 10, 5)

  defp float_to_string(<<bits::8-native-unsigned>>, :f, 8),
    do: Nx.Ryu.bits_to_decimal(bits, 2, 5)

  defp float_to_string(<<bits::16-native-unsigned>>, :bf, 16),
    do: Nx.Ryu.bits_to_decimal(bits, 7, 8)

  defp float_to_string(<<bits::8-native-unsigned>>, :f8_e4m3fn, 8),
    do: Nx.Ryu.bits_to_decimal(bits, 3, 4, :fn)

  defp complex_to_string(re, im, float_unit) do
    re_str = float_to_string(re, :f, float_unit)
    im_str = float_to_string(im, :f, float_unit)

    im_str =
      case im_str do
        "-" <> _ -> im_str
        s -> "+" <> s
      end

    re_str <> im_str <> "i"
  end
end