From 05d66a18d7f92b54c7e564490afd0ee5c353cf1c Mon Sep 17 00:00:00 2001 From: Chapaman <14204271+Chapaman@users.noreply.github.com> Date: Wed, 1 Jul 2026 22:23:56 -0300 Subject: [PATCH 1/6] Document backend behaviour with markdown guides instead of ghost modules --- exla/guides/backend_documentation/index.md | 24 ++ exla/guides/backend_documentation/nx.md | 264 ++++++++++++++++++ .../backend_documentation/nx_lin_alg.md | 216 ++++++++++++++ exla/lib/exla.ex | 7 + exla/mix.exs | 20 +- exla/test/exla/backend_documentation_test.exs | 7 + nx/guides/backend_documentation/convention.md | 54 ++++ nx/lib/nx.ex | 112 ++++++++ nx/lib/nx/backend.ex | 4 + nx/lib/nx/lin_alg.ex | 40 +++ nx/mix.exs | 7 +- nx/test/nx/backend_documentation_test.exs | 5 + torchx/guides/backend_documentation/index.md | 24 ++ torchx/guides/backend_documentation/nx.md | 221 +++++++++++++++ .../backend_documentation/nx_lin_alg.md | 199 +++++++++++++ torchx/lib/torchx.ex | 8 + torchx/lib/torchx/backend.ex | 5 + torchx/mix.exs | 19 +- .../torchx/backend_documentation_test.exs | 7 + 19 files changed, 1238 insertions(+), 5 deletions(-) create mode 100644 exla/guides/backend_documentation/index.md create mode 100644 exla/guides/backend_documentation/nx.md create mode 100644 exla/guides/backend_documentation/nx_lin_alg.md create mode 100644 exla/test/exla/backend_documentation_test.exs create mode 100644 nx/guides/backend_documentation/convention.md create mode 100644 nx/test/nx/backend_documentation_test.exs create mode 100644 torchx/guides/backend_documentation/index.md create mode 100644 torchx/guides/backend_documentation/nx.md create mode 100644 torchx/guides/backend_documentation/nx_lin_alg.md create mode 100644 torchx/test/torchx/backend_documentation_test.exs diff --git a/exla/guides/backend_documentation/index.md b/exla/guides/backend_documentation/index.md new file mode 100644 index 0000000000..32fb157aa3 --- /dev/null +++ b/exla/guides/backend_documentation/index.md @@ -0,0 +1,24 @@ +# Backend documentation + +EXLA-specific documentation for Nx backend behaviour. + +These guides describe how EXLA lowers and executes Nx operations. They mirror +the structure of the Nx API (see the [backend documentation convention](https://hexdocs.pm/nx/backend_documentation-convention.html)) +but are **not callable** — use `Nx` and `compiler: EXLA` in your code. + +## Guides + + * [Nx](backend_documentation-nx.html) — top-level `Nx` blocks, transfers, and `defn` integration + * [Nx.LinAlg](backend_documentation-nx_lin_alg.html) — linear algebra blocks and related lowerings + +Implementation code lives in `EXLA.Defn`, `EXLA.CustomCall`, and related modules. + +## Example + +```elixir +iex> Nx.take(Nx.tensor([10, 20, 30]), Nx.tensor([0, 2])) +#Nx.Tensor< + s32[2] + [10, 30] +> +``` diff --git a/exla/guides/backend_documentation/nx.md b/exla/guides/backend_documentation/nx.md new file mode 100644 index 0000000000..d510b06945 --- /dev/null +++ b/exla/guides/backend_documentation/nx.md @@ -0,0 +1,264 @@ +# Nx (EXLA) + +```elixir +iex> Nx.take(Nx.tensor([10, 20, 30]), Nx.tensor([0, 2])) +#Nx.Tensor< + s32[2] + [10, 30] +> +``` + +EXLA implementation notes for top-level `Nx` operations. + +Each function below documents how EXLA handles the corresponding `Nx` API when +`compiler: EXLA` is used (or when `EXLA.Backend` stores tensors). + +Linear algebra is documented separately in [Nx.LinAlg](backend_documentation-nx_lin_alg.html). + +Many `Nx` functions delegate to `Nx.block/4`. During EXLA compilation, +specialized lowerings exist for some block tags; otherwise EXLA compiles the +default callback as an XLA subcomputation. + +## take/3 + +Tensor indexing (`%Nx.Block.Take{}`). + +### Lowering + +Lowered to StableHLO `gather` in `EXLA.Defn` (not a custom call). + +### Options + +* `:axis` — honoured; indices are gathered along this axis + +### Platforms + +All EXLA clients. + +## take_along_axis/3 + +Take along axis (`%Nx.Block.TakeAlongAxis{}`). + +### Lowering + +No dedicated gather lowering — EXLA compiles the default Nx callback as an XLA +subcomputation. + +### Options + +* `:axis` — honoured by the default callback + +## top_k/2 + +Top-k values and indices (`%Nx.Block.TopK{}`). + +### Lowering + +Lowered to StableHLO `top_k` via `EXLA.MLIR.Value.top_k/3`. + +### Options + +* `:k` — number of elements to return per slice + +## fft2/2 + +Two-dimensional FFT (`%Nx.Block.FFT2{}`). + +### Lowering + +Lowered to StableHLO FFT ops via `EXLA.MLIR.Value.fft/4` with mode `:fft`. + +### Options + +* `:lengths`, `:axes`, `:eps` — forwarded to the lowering; `:eps` may be used +for cleanup in the default callback when compilation falls back + +## ifft2/2 + +Two-dimensional inverse FFT (`%Nx.Block.IFFT2{}`). + +### Lowering + +Lowered to StableHLO FFT ops with mode `:ifft`. + +## rfft/2 + +Real FFT (`%Nx.Block.RFFT{}`). + +### Lowering + +Lowered to StableHLO real FFT via `Value.fft/4` with mode `:rfft`. Input is +treated as real; output type is complex per the expression template. + +## irfft/2 + +Inverse real FFT (`%Nx.Block.IRFFT{}`). + +### Lowering + +Lowered to StableHLO IRFFT via `Value.fft/4` with mode `:irfft`. + +## cumulative_sum/2 + +Cumulative sum (`%Nx.Block.CumulativeSum{}`). + +### Lowering + +Default callback compiled as an XLA subcomputation (no native custom call). + +### Options + +* `:axis`, `:reverse` — honoured by the default callback + +## cumulative_product/2 + +Cumulative product (`%Nx.Block.CumulativeProduct{}`). + +### Lowering + +Default callback compiled as an XLA subcomputation. + +## cumulative_min/2 + +Cumulative minimum (`%Nx.Block.CumulativeMin{}`). + +### Lowering + +Default callback compiled as an XLA subcomputation. + +## cumulative_max/2 + +Cumulative maximum (`%Nx.Block.CumulativeMax{}`). + +### Lowering + +Default callback compiled as an XLA subcomputation. + +## all_close/3 + +All-close comparison (`%Nx.Block.AllClose{}`). + +### Lowering + +Default callback compiled as an XLA subcomputation. + +### Options + +* `:rtol`, `:atol`, `:equal_nan` — honoured by the default callback + +## logical_not/1 + +Logical not (`%Nx.Block.LogicalNot{}`). + +### Lowering + +Default callback compiled as an XLA subcomputation (delegates to element-wise +equality in the reference implementation). + +## phase/1 + +Complex phase (`%Nx.Block.Phase{}`). + +### Lowering + +Default callback compiled as an XLA subcomputation. + +## runtime_call/4 + +Runtime Elixir callback from `defn` (`runtime_call` expression). + +### Lowering + +On **host** and **CUDA** clients, EXLA emits a device-side callback bridged +through `EXLA.MLIR.Value.runtime_call/3`. Callback tensors are materialized for +the Elixir function; results must match the output template. + +### Platforms + +* **Host / CUDA** — supported +* **ROCm / TPU** — raises with a message to use `:host` or `:cuda` + +### Outside `defn` + +Executes the callback directly on the input backend without compilation. + +### Warnings + +Avoid `Nx.backend_transfer/2` on callback tensors inside the function when +using `Nx.Defn.Evaluator`. Avoid running other Nx computations on the same GPU +device from within the callback (deadlock risk). + +## backend_copy/2 + +Copy tensor data to another backend (`EXLA.Backend.backend_copy/3`). + +### Behaviour + +Copying to `EXLA.Backend` on the same client and device returns the tensor +unchanged. Cross-device copies use `EXLA.DeviceBuffer.copy_to_device/3`. +Copying out reads device memory via `EXLA.DeviceBuffer.read/1` and delegates +to the target backend's `from_binary/3`. + +### Options + +* `:client`, `:device_id` — select the EXLA client and device + +## backend_transfer/2 + +Transfer tensor data to another backend (`EXLA.Backend.backend_transfer/3`). + +### Behaviour + +Same as `backend_copy/2` followed by deallocation of the source device buffer +when leaving `EXLA.Backend`. + +## backend_deallocate/1 + +Deallocate device memory (`EXLA.Backend.backend_deallocate/1`). + +### Behaviour + +Calls `EXLA.DeviceBuffer.deallocate/1` for tensors on `EXLA.Backend`. + +## to_pointer/2 + +Export device memory as a pointer (`EXLA.Backend.to_pointer/2`). + +### Modes + +* `:local` — supported on **host** and **CUDA** clients; returns a local +address integer +* `:ipc` — supported on **host** (`shm_open` handle) and **CUDA** (IPC +handle plus device id) + +### Options + +* `:permissions` — octal file mode for host IPC shared memory (default +`0o400`) + +### Limitations + +Not supported for ROCm, TPU, or non-device buffers. + +## from_pointer/5 + +Import device memory from a pointer (`EXLA.Backend.from_pointer/5`). + +### Behaviour + +Creates an `EXLA.DeviceBuffer` from a local address or IPC handle on **host** +or **CUDA** clients. Pointer `data_size` must match the tensor byte size. + +### Options + +* `:client`, `:device_id` — destination device +* `:names` — tensor names for the result template + +## to_batched/3 + +Stream tensors in batches (`EXLA.Backend.to_batched/3`). + +### Behaviour + +Splits the leading axis via XLA slice and concatenate operations on device +buffers. Supports `:leftover` `:repeat` and `:discard` like `Nx.to_batched/2`. diff --git a/exla/guides/backend_documentation/nx_lin_alg.md b/exla/guides/backend_documentation/nx_lin_alg.md new file mode 100644 index 0000000000..e61a16e5fb --- /dev/null +++ b/exla/guides/backend_documentation/nx_lin_alg.md @@ -0,0 +1,216 @@ +# Nx.LinAlg (EXLA) + +```elixir +iex> a = Nx.tensor([[4.0, 2.0], [2.0, 3.0]]) +iex> l = Nx.LinAlg.cholesky(a) +iex> Nx.shape(l) +{2, 2} +``` + +EXLA implementation notes for `Nx.LinAlg`. + +Each function below documents how EXLA handles the corresponding `Nx.LinAlg` +operation when `compiler: EXLA` is used. + +Operations that use `Nx.block/4` are dispatched through `EXLA.Defn`. When +`EXLA.CustomCall.call/4` returns `:skip`, EXLA compiles the default block +callback as an XLA subcomputation. When it returns `{:ok, spec}`, EXLA emits a +StableHLO `custom_call` instead. + +On GPU and TPU clients, the `:precision` option passed to `Nx.Defn.jit/2` +affects accumulator precision for many composed linear algebra graphs compiled +from the default callbacks. For decomposition verification tests, prefer +`precision: :highest`. + +## cholesky/1 + +Cholesky decomposition (`%Nx.Block.LinAlg.Cholesky{}`). + +### Lowering + +EXLA does **not** provide a native custom call for Cholesky. The default +callback (`Nx.LinAlg.Cholesky.cholesky/1`) is compiled as an XLA +subcomputation. + +### Supported platforms + +All EXLA clients (`:host`, `:cuda`, `:rocm`, `:tpu`) use the compiled default +implementation. + +### Types + +Follows Nx promotion rules: integer and unsigned inputs are promoted to +floating point in the public API before the block is invoked. + +### Numerical notes + +Positive-definiteness is not validated at compile time. Ill-conditioned inputs +may produce `NaN` values without raising, consistent with the reference Nx +implementation. + +## solve/2 + +Linear system solve (`%Nx.Block.LinAlg.Solve{}`). + +### Lowering + +EXLA compiles the default callback, which expresses `solve/2` as an LU +factorization followed by triangular solves. + +### Supported platforms + +All EXLA clients use the compiled default implementation. + +### Types + +Output type is floating point per `Nx.LinAlg.solve/2`. The solve block depends +on `lu/1` and `triangular_solve/3` inside the compiled graph. + +## qr/2 + +QR decomposition (`%Nx.Block.LinAlg.QR{}`). + +### Lowering + +On the **host** client, with **real** input element types, EXLA emits a +native CPU custom call via `EXLA.CustomCall` when the dtype is supported: + +* `{:f, 32}` → `qr_cpu_custom_call_f32` +* `{:f, 64}` → `qr_cpu_custom_call_f64` +* `{:f, 16}` → `qr_cpu_custom_call_f16` +* `{:bf, 16}` → `qr_cpu_custom_call_bf16` +* `{:s, _}` and `{:u, _}` → `qr_cpu_custom_call_f32` with operand cast to +`f32` + +Otherwise EXLA compiles the default callback (`Nx.LinAlg.QR.qr/2`). + +### Supported platforms + +* **Host** — native custom call for supported real dtypes; default callback +for complex inputs and unsupported types +* **CUDA / ROCm / TPU** — default callback only (no native QR custom call) + +### Options + +* `:mode` — honoured by whichever lowering is selected (`:reduced` or +`:complete`) +* `:eps` — used by the default callback only; ignored by the native custom +call + +### Numerical notes + +Native and default lowerings may differ slightly in orthogonality of `q` and +triangular structure of `r` within floating-point tolerance. Reconstruction +`q · r` should match the input within reasonable tolerances for the dtype. + +## eigh/2 + +Hermitian eigendecomposition (`%Nx.Block.LinAlg.Eigh{}`). + +### Lowering + +On the **host** client, with **real** input element types, EXLA emits a +native CPU custom call when the dtype is supported: + +* `{:f, 32}` → `eigh_cpu_custom_call_f32` +* `{:f, 64}` → `eigh_cpu_custom_call_f64` +* `{:s, _}` and `{:u, _}` → `eigh_cpu_custom_call_f32` with operand cast to +`f32` + +Otherwise EXLA compiles the default callback (`Nx.LinAlg.BlockEigh.eigh/2` or +`Nx.LinAlg.Eigh.eigh/2` depending on input rank). + +### Supported platforms + +* **Host** — native custom call for supported real dtypes; default callback +for complex inputs and unsupported types +* **CUDA / ROCm / TPU** — default callback only + +### Options + +* `:max_iter` and `:eps` — honoured by the default callback; ignored by the +native custom call + +### Numerical notes + +Eigenvectors are not unique (sign and degenerate subspaces). Compare results +using reconstruction `v · diag(λ) · vᵀ` rather than direct eigenvector equality +across backends. + +## svd/2 + +Singular value decomposition (`%Nx.Block.LinAlg.SVD{}`). + +### Lowering + +EXLA compiles the default callback (`Nx.LinAlg.SVD.svd/2`). There is no native +SVD custom call. + +### Supported platforms + +All EXLA clients use the compiled default implementation. + +### Options + +* `:max_iter` — honoured (default iterative algorithm in Nx) +* `:full_matrices?` — honoured + +### Numerical notes + +SVD is iterative; results may vary across devices and precision settings. +Singular values are the most stable quantity to compare across backends. +Use `precision: :highest` on GPU when verifying decompositions. + +## lu/1 + +LU decomposition with partial pivoting (`%Nx.Block.LinAlg.LU{}`). + +### Lowering + +EXLA compiles the default callback (`Nx.LinAlg.LU.lu/1`) as an XLA +subcomputation. Permutation matrix `p` uses integer (`s32`) type; `l` and `u` +are floating point. + +### Supported platforms + +All EXLA clients use the compiled default implementation. + +### Numerical notes + +For rank-deficient inputs, `u` may contain zeros on the diagonal. This matches +the reference Nx behaviour. + +## determinant/1 + +Matrix determinant (`%Nx.Block.LinAlg.Determinant{}`). + +### Lowering + +EXLA compiles the default callback. Small matrices use closed-form expressions; +larger matrices use LU-based determinant computation inside the compiled graph. + +### Supported platforms + +All EXLA clients use the compiled default implementation. + +### Types + +Output is floating point. Complex determinants are supported when the input is +complex. + +## triangular_solve/3 + +Triangular solve (direct `triangular_solve` callback, not an `Nx.block/4`). + +### Lowering + +Lowered to the StableHLO `triangular_solve` op via `EXLA.MLIR.Value.triangular_solve/4`. + +### Supported platforms + +All EXLA clients. + +### Options + +EXLA honours `lower:`, `left_side:`, and `transform_a:` as defined in +`Nx.LinAlg.triangular_solve/3`. diff --git a/exla/lib/exla.ex b/exla/lib/exla.ex index 9be5dc13db..cb8d162e50 100644 --- a/exla/lib/exla.ex +++ b/exla/lib/exla.ex @@ -82,6 +82,13 @@ defmodule EXLA do your block tag struct; see `EXLA.CustomCall` for the `call/4` contract, including returning `:skip` to fall back to the block's default Elixir callback. + Per-operation lowering notes for Nx functions are documented in the + [Backend documentation](backend_documentation.html) guides (for example + [Nx](backend_documentation-nx.html) and + [Nx.LinAlg](backend_documentation-nx_lin_alg.html)). See also the + [backend documentation convention](https://hexdocs.pm/nx/backend_documentation-convention.html) + in the Nx package. + ## Clients The `EXLA` library uses a client for compiling and executing code. diff --git a/exla/mix.exs b/exla/mix.exs index 035450b640..9dc28f4cbe 100644 --- a/exla/mix.exs +++ b/exla/mix.exs @@ -70,7 +70,8 @@ defmodule EXLA.MixProject do defp deps do [ - {:nx, path: "../nx"}, + {:nx, "~> 0.12.0"}, + # {:nx, path: "../nx"}, {:telemetry, "~> 0.4.0 or ~> 1.0"}, {:xla, "~> 0.10.0", runtime: false}, {:fine, "~> 0.1", runtime: false}, @@ -87,9 +88,24 @@ defmodule EXLA.MixProject do source_url_pattern: "#{@source_url}/blob/v#{@version}/exla/%{path}#L%{line}", extras: [ "guides/rotating-image.livemd", - "CHANGELOG.md" + "CHANGELOG.md", + "guides/backend_documentation/index.md": [ + filename: "backend_documentation", + title: "Backend documentation" + ], + "guides/backend_documentation/nx.md": [ + filename: "backend_documentation-nx", + title: "Nx (EXLA)" + ], + "guides/backend_documentation/nx_lin_alg.md": [ + filename: "backend_documentation-nx_lin_alg", + title: "Nx.LinAlg (EXLA)" + ] ], skip_undefined_reference_warnings_on: ["CHANGELOG.md"], + groups_for_extras: [ + "Backend documentation": ~r"^guides/backend_documentation/" + ], groups_for_modules: [ # EXLA, # EXLA.Backend, diff --git a/exla/test/exla/backend_documentation_test.exs b/exla/test/exla/backend_documentation_test.exs new file mode 100644 index 0000000000..b6de7ca3cd --- /dev/null +++ b/exla/test/exla/backend_documentation_test.exs @@ -0,0 +1,7 @@ +defmodule EXLA.BackendDocumentationTest do + use ExUnit.Case, async: true + + doctest_file "guides/backend_documentation/index.md" + doctest_file "guides/backend_documentation/nx.md" + doctest_file "guides/backend_documentation/nx_lin_alg.md" +end diff --git a/nx/guides/backend_documentation/convention.md b/nx/guides/backend_documentation/convention.md new file mode 100644 index 0000000000..c9df6ffea3 --- /dev/null +++ b/nx/guides/backend_documentation/convention.md @@ -0,0 +1,54 @@ +# Backend documentation convention + +Nx exposes a stable public API (`Nx`, `Nx.LinAlg`, and so on) while individual +backends implement the underlying callbacks and optional `Nx.block/4` lowerings. +Users who need to know *how* a specific backend executes an operation should +not have to read backend source code. + +## Documentation guides + +Each backend may publish **backend documentation guides** that mirror the Nx +module structure and document backend-specific behaviour for each callback. +The naming convention is: + + guides/backend_documentation/nx.md # mirrors top-level `Nx` + guides/backend_documentation/nx_lin_alg.md # mirrors `Nx.LinAlg` + +For example, EXLA documents operations in its [Backend documentation](https://hexdocs.pm/exla/backend_documentation.html) +guides, and Torchx in its [Backend documentation](https://hexdocs.pm/torchx/backend_documentation.html) +guides. + +These guides describe primitives that backends implement, without being part of the +callable Nx API. + +## What to document + +Backends should document, for each mirrored function: + + * whether the operation is lowered natively, compiled from the default + `Nx.block/4` callback, or handled through another mechanism (for example + a direct StableHLO op) + * supported types, devices, and platforms + * options that the backend honours, ignores, or overrides + * numerical behaviour, performance trade-offs, and known limitations relative + to the reference implementation in Nx + +Cross-link from the Nx function doc when appropriate (for example +`See EXLA and Torchx backend documentation for qr/2`). + +## Relation to `Nx.Backend` + +Most tensor operations go through `Nx.Backend` callbacks. Operations that need +a portable default with optional native acceleration use `Nx.block/4` and the +`Nx.Block.*` tags — see `Nx.Backend.block/4`. Backend documentation guides +should cover both callback implementations and block lowerings. + +## Example + +```elixir +iex> Nx.take(Nx.tensor([10, 20, 30]), Nx.tensor([0, 2])) +#Nx.Tensor< + s32[2] + [10, 30] +> +``` diff --git a/nx/lib/nx.ex b/nx/lib/nx.ex index 77dd527889..f6a8fb6c10 100644 --- a/nx/lib/nx.ex +++ b/nx/lib/nx.ex @@ -39,6 +39,13 @@ defmodule Nx do * `Nx.Constants` declares many constants commonly used in numerical code + Backend-specific notes for extensible operations (`Nx.block/4`, transfers, and + related APIs) are documented in the [backend documentation convention](backend_documentation-convention.html) + guide and in each backend's **Backend documentation** pages on HexDocs (for example + [EXLA](https://hexdocs.pm/exla/backend_documentation.html) and + [Torchx](https://hexdocs.pm/torchx/backend_documentation.html)). + + Continue reading this documentation for an overview of creating, broadcasting, and accessing/slicing Nx tensors. @@ -1615,6 +1622,11 @@ defmodule Nx do iex> Nx.take_diagonal(Nx.iota({3, 3}), offset: -4) ** (ArgumentError) absolute value of offset must be less than length of axis 0 when negative, got: -4 + + ## Backend notes + + See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#take-3) and + [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#take-3). """ @doc type: :creation def take_diagonal(tensor, opts \\ []) do @@ -2271,6 +2283,11 @@ defmodule Nx do Inside `defn`, this builds an expression node that is supported by compilers. Outside of `defn` or in backends without special support, it executes `fun` directly and validates that the result matches the provided template. + + ## Backend notes + + See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#runtime_call-4) and + [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#runtime_call-4). """ @doc type: :backend def runtime_call(output, tensor_or_container, opts \\ [], fun) when is_function(fun, 2) do @@ -2515,6 +2532,11 @@ defmodule Nx do s32 [3, 1] > + + ## Backend notes + + See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#to_batched-3) and + [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#to_batched-3). """ @doc type: :conversion def to_batched(tensor, batch_size, opts \\ []) @@ -4980,6 +5002,11 @@ defmodule Nx do ] > + + ## Backend notes + + See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#backend_copy-2) and + [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#backend_copy-2). """ @doc type: :backend def backend_copy(tensor_or_container, backend \\ Nx.BinaryBackend) do @@ -5030,6 +5057,11 @@ defmodule Nx do tensor = Nx.backend_transfer(device_tensor) + + ## Backend notes + + See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#backend_transfer-2) and + [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#backend_transfer-2). """ @doc type: :backend def backend_transfer(tensor_or_container, backend \\ Nx.BinaryBackend) do @@ -5049,6 +5081,11 @@ defmodule Nx do It returns either `:ok` or `:already_deallocated`. Note: This function cannot be used in `defn`. + + ## Backend notes + + See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#backend_deallocate-1) and + [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#backend_deallocate-1). """ @doc type: :backend def backend_deallocate(tensor_or_container) do @@ -7223,6 +7260,11 @@ defmodule Nx do [0, 1, 0] > + + ## Backend notes + + See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#logical_not-1) and + [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#logical_not-1). """ @doc type: :element def logical_not(tensor) do @@ -8690,6 +8732,11 @@ defmodule Nx do f32[2] [1.1071488, 2.6779451] > + + ## Backend notes + + See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#phase-1) and + [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#phase-1). """ @doc type: :element def phase(tensor) do @@ -9276,6 +9323,11 @@ defmodule Nx do [0, 0] ] > + + ## Backend notes + + See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#all_close-3) and + [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#all_close-3). """ @doc type: :aggregation def all_close(a, b, opts \\ []) do @@ -11680,6 +11732,11 @@ defmodule Nx do [2, 3, 6] ] > + + ## Backend notes + + See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#cumulative_sum-2) and + [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#cumulative_sum-2). """ @doc type: :cumulative def cumulative_sum(tensor, opts \\ []), @@ -11756,6 +11813,11 @@ defmodule Nx do [2, 2, 6] ] > + + ## Backend notes + + See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#cumulative_product-2) and + [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#cumulative_product-2). """ @doc type: :cumulative def cumulative_product(tensor, opts \\ []), @@ -11832,6 +11894,11 @@ defmodule Nx do [2, 1, 1] ] > + + ## Backend notes + + See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#cumulative_min-2) and + [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#cumulative_min-2). """ @doc type: :cumulative def cumulative_min(tensor, opts \\ []), @@ -11908,6 +11975,11 @@ defmodule Nx do [2, 2, 3] ] > + + ## Backend notes + + See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#cumulative_max-2) and + [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#cumulative_max-2). """ @doc type: :cumulative def cumulative_max(tensor, opts \\ []), @@ -14840,6 +14912,11 @@ defmodule Nx do iex> idx = Nx.tensor([[2.0], [1.0], [2.0]], type: :f32) iex> Nx.take_along_axis(tensor, idx, axis: 1) ** (ArgumentError) indices must be an integer tensor, got {:f, 32} + + ## Backend notes + + See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#take_along_axis-3) and + [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#take_along_axis-3). """ @doc type: :indexed def take_along_axis(tensor, indices, opts \\ []) when is_list(opts) do @@ -15667,6 +15744,11 @@ defmodule Nx do iex> Nx.top_k(a, k: 1) ** (ArgumentError) top_k input must have at least rank 1 + + ## Backend notes + + See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#top_k-2) and + [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#top_k-2). """ @doc type: :ndim def top_k(tensor, opts \\ []) do @@ -16902,6 +16984,11 @@ defmodule Nx do iex> Nx.fft2(Nx.tensor([1, 1]), length: :invalid) ** (ArgumentError) expected a tensor with rank > 1, got tensor with rank 1 + + ## Backend notes + + See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#fft2-2) and + [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#fft2-2). """ @doc type: :ndim def fft2(tensor, opts \\ []), do: call_fft2(tensor, opts, :fft2) @@ -17060,6 +17147,11 @@ defmodule Nx do iex> Nx.ifft2(Nx.tensor([1, 1]), length: :invalid) ** (ArgumentError) expected a tensor with rank > 1, got tensor with rank 1 + + ## Backend notes + + See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#ifft2-2) and + [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#ifft2-2). """ @doc type: :ndim def ifft2(tensor, opts \\ []), do: call_fft2(tensor, opts, :ifft2) @@ -17234,6 +17326,11 @@ defmodule Nx do iex> Nx.rfft(Nx.tensor([1.0, 1.0]), length: :invalid) ** (ArgumentError) expected an integer or :power_of_two as length, got: :invalid + + ## Backend notes + + See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#rfft-2) and + [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#rfft-2). """ @doc type: :ndim def rfft(tensor, opts \\ []) do @@ -17355,6 +17452,11 @@ defmodule Nx do iex> Nx.irfft(Nx.tensor([1.0, 1.0]), length: :invalid) ** (ArgumentError) expected a positive integer as length, got: :invalid + + ## Backend notes + + See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#irfft-2) and + [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#irfft-2). """ @doc type: :ndim def irfft(tensor, opts \\ []) do @@ -17632,6 +17734,11 @@ defmodule Nx do #=> [10, 20, 30] #=> ] #=> > + + ## Backend notes + + See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#from_pointer-5) and + [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#from_pointer-5). """ @doc type: :creation def from_pointer(backend, pointer, type, shape, opts \\ []) @@ -17677,6 +17784,11 @@ defmodule Nx do t = Nx.s32([1, 2, 3]) Nx.to_pointer(t, mode: :ipc) #=> %Nx.Pointer{kind: :ipc, address: nil, data_size: 32, handle: "some-ipc-handle"} + + ## Backend notes + + See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#to_pointer-2) and + [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#to_pointer-2). """ @doc type: :creation def to_pointer(tensor, opts \\ []) do diff --git a/nx/lib/nx/backend.ex b/nx/lib/nx/backend.ex index 646356228a..bc91bbdd02 100644 --- a/nx/lib/nx/backend.ex +++ b/nx/lib/nx/backend.ex @@ -28,6 +28,10 @@ defmodule Nx.Backend do 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 implementation specifics — see the + [backend documentation convention](backend_documentation-convention.html) guide. """ @type t :: %{__struct__: atom()} diff --git a/nx/lib/nx/lin_alg.ex b/nx/lib/nx/lin_alg.ex index e7e166c6ab..ee9e7c00d3 100644 --- a/nx/lib/nx/lin_alg.ex +++ b/nx/lib/nx/lin_alg.ex @@ -136,6 +136,11 @@ defmodule Nx.LinAlg do ] ] > + + ## Backend notes + + See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx_lin_alg.html#cholesky-1) and + [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx_lin_alg.html#cholesky-1). """ def cholesky(tensor) do %T{vectorized_axes: vectorized_axes} = tensor = Nx.to_tensor(tensor) @@ -584,6 +589,11 @@ defmodule Nx.LinAlg do iex> Nx.LinAlg.triangular_solve(a, Nx.tensor([1, 2, 1]), transform_a: :other) ** (ArgumentError) invalid value for :transform_a option, expected :none, :transpose, or :conjugate, got: :other + + ## Backend notes + + See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx_lin_alg.html#triangular_solve-3) and + [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx_lin_alg.html#triangular_solve-3). """ def triangular_solve(a, b, opts \\ []) do opts = keyword!(opts, lower: true, left_side: true, transform_a: :none) @@ -697,6 +707,11 @@ defmodule Nx.LinAlg do iex> Nx.LinAlg.solve(Nx.tensor([[3, 0, 0, 0], [2, 1, 0, 0], [1, 1, 1, 1]]), Nx.tensor([4])) ** (ArgumentError) `a` tensor has incompatible dimensions, expected a square matrix or a batch of square matrices, got: {3, 4} + + ## Backend notes + + See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx_lin_alg.html#solve-2) and + [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx_lin_alg.html#solve-2). """ # IMPORTANT: This function cannot be a defn because # optional needs to work on the actual backend. @@ -1123,6 +1138,11 @@ defmodule Nx.LinAlg do iex> t = Nx.tensor([[-3, 2, 1], [0, 1, 1], [0, 0, -1]]) iex> Nx.LinAlg.qr(t, mode: :error_test) ** (ArgumentError) invalid :mode received. Expected one of [:reduced, :complete], received: :error_test + + ## Backend notes + + See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx_lin_alg.html#qr-2) and + [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx_lin_alg.html#qr-2). """ def qr(tensor, opts \\ []) do opts = keyword!(opts, mode: :reduced, eps: 1.0e-10) @@ -1383,6 +1403,11 @@ defmodule Nx.LinAlg do iex> Nx.LinAlg.eigh(Nx.tensor([[1, 2, 3], [4, 5, 6]])) ** (ArgumentError) tensor must be a square matrix or a batch of square matrices, got shape: {2, 3} + + ## Backend notes + + See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx_lin_alg.html#eigh-2) and + [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx_lin_alg.html#eigh-2). """ def eigh(tensor, opts \\ []) do opts = keyword!(opts, max_iter: 1_000, eps: 1.0e-4) @@ -1506,6 +1531,11 @@ defmodule Nx.LinAlg do [0.0, 0.0, 1.0] ] > + + ## Backend notes + + See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx_lin_alg.html#svd-2) and + [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx_lin_alg.html#svd-2). """ def svd(tensor, opts \\ []) do opts = keyword!(opts, max_iter: 100, full_matrices?: true) @@ -1736,6 +1766,11 @@ defmodule Nx.LinAlg do iex> Nx.LinAlg.lu(Nx.tensor([[1, 1, 1, 1], [-1, 4, 4, -1], [4, -2, 2, 0]])) ** (ArgumentError) tensor must be a square matrix or a batch of square matrices, got shape: {3, 4} + + ## Backend notes + + See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx_lin_alg.html#lu-1) and + [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx_lin_alg.html#lu-1). """ def lu(tensor) do %T{vectorized_axes: vectorized_axes} = tensor = Nx.to_tensor(tensor) @@ -1985,6 +2020,11 @@ defmodule Nx.LinAlg do -0.0-6.0i > + + ## Backend notes + + See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx_lin_alg.html#determinant-1) and + [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx_lin_alg.html#determinant-1). """ # IMPORTANT: This function cannot be a defn because # optional needs to work on the actual backend. diff --git a/nx/mix.exs b/nx/mix.exs index d148634b8e..df9585e892 100644 --- a/nx/mix.exs +++ b/nx/mix.exs @@ -74,7 +74,11 @@ defmodule Nx.MixProject do "guides/advanced/aggregation.livemd", "guides/advanced/automatic_differentiation.livemd", "guides/advanced/complex_fft.livemd", - "guides/exercises/exercises-1-20.livemd" + "guides/exercises/exercises-1-20.livemd", + "guides/backend_documentation/convention.md": [ + filename: "backend_documentation-convention", + title: "Backend documentation convention" + ] ], skip_undefined_reference_warnings_on: ["CHANGELOG.md"], groups_for_docs: [ @@ -132,6 +136,7 @@ defmodule Nx.MixProject do groups_for_extras: [ "Getting Started": ~r"^guides/getting_started/", Cheatsheets: ~r"^guides/cheatsheets/", + "Backend documentation": ~r"^guides/backend_documentation/", Exercises: ~r"^guides/exercises/", Advanced: ~r"^guides/advanced/" ] diff --git a/nx/test/nx/backend_documentation_test.exs b/nx/test/nx/backend_documentation_test.exs new file mode 100644 index 0000000000..3dfc6d4ae4 --- /dev/null +++ b/nx/test/nx/backend_documentation_test.exs @@ -0,0 +1,5 @@ +defmodule Nx.BackendDocumentationTest do + use ExUnit.Case, async: true + + doctest_file "guides/backend_documentation/convention.md" +end diff --git a/torchx/guides/backend_documentation/index.md b/torchx/guides/backend_documentation/index.md new file mode 100644 index 0000000000..c056db57b0 --- /dev/null +++ b/torchx/guides/backend_documentation/index.md @@ -0,0 +1,24 @@ +# Backend documentation + +Torchx-specific documentation for Nx backend behaviour. + +These guides describe how Torchx executes Nx operations through LibTorch. +They mirror the structure of the Nx API (see the [backend documentation convention](https://hexdocs.pm/nx/backend_documentation-convention.html)) +but are **not callable** — use `Nx` with `Torchx.Backend` in your code. + +## Guides + + * [Nx](backend_documentation-nx.html) — top-level `Nx` blocks, transfers, and callbacks + * [Nx.LinAlg](backend_documentation-nx_lin_alg.html) — linear algebra blocks and related callbacks + +Implementation code lives in `Torchx.Backend` and the `Torchx` NIF bindings. + +## Example + +```elixir +iex> Nx.take(Nx.tensor([10, 20, 30]), Nx.tensor([0, 2])) +#Nx.Tensor< + s32[2] + [10, 30] +> +``` diff --git a/torchx/guides/backend_documentation/nx.md b/torchx/guides/backend_documentation/nx.md new file mode 100644 index 0000000000..9998d71b90 --- /dev/null +++ b/torchx/guides/backend_documentation/nx.md @@ -0,0 +1,221 @@ +# Nx (Torchx) + +```elixir +iex> Nx.take(Nx.tensor([10, 20, 30]), Nx.tensor([0, 2])) +#Nx.Tensor< + s32[2] + [10, 30] +> +``` + +Torchx implementation notes for top-level `Nx` operations. + +Each function below documents how Torchx handles the corresponding `Nx` API +when `Torchx.Backend` is active. + +Linear algebra is documented separately in [Nx.LinAlg](backend_documentation-nx_lin_alg.html). + +Torchx is an **Nx backend**, not a `defn` compiler. Block lowerings run through +`Torchx.Backend.block/4` during eager execution or when another compiler +(such as `Nx.Defn.Evaluator`) invokes backend callbacks. + +## take/3 + +Tensor indexing (`%Nx.Block.Take{}`). + +### Lowering + +No native LibTorch path — Torchx invokes the default Nx block callback, which +composes `Nx` gather operations on the Torchx backend. + +### Options + +* `:axis` — honoured by the default callback + +## take_along_axis/3 + +Take along axis (`%Nx.Block.TakeAlongAxis{}`). + +### Lowering + +Native `Torchx.gather/2` along the configured axis (LibTorch `gather`). + +## top_k/2 + +Top-k values and indices (`%Nx.Block.TopK{}`). + +### Lowering + +Native `Torchx.top_k/2`. + +### Options + +* `:k` — passed to LibTorch + +## fft2/2 + +Two-dimensional FFT (`%Nx.Block.FFT2{}`). + +### Lowering + +Native `Torchx.fft2/3`. + +### Numerical notes + +Signed zeros and `inspect` formatting may differ from `Nx.BinaryBackend`; see +`Torchx.NxBlockTest` for numerical comparisons. + +## ifft2/2 + +Two-dimensional inverse FFT (`%Nx.Block.IFFT2{}`). + +### Lowering + +Native `Torchx.ifft2/3`. + +## rfft/2 + +Real FFT (`%Nx.Block.RFFT{}`). + +### Lowering + +Native `Torchx.rfft/3`. + +## irfft/2 + +Inverse real FFT (`%Nx.Block.IRFFT{}`). + +### Lowering + +Native `Torchx.irfft/3`. + +## cumulative_sum/2 + +Cumulative sum (`%Nx.Block.CumulativeSum{}`). + +### Lowering + +Native `Torchx.cumulative_sum/2` with optional `flip` when `:reverse` is true. + +### Options + +* `:axis`, `:reverse` — honoured + +## cumulative_product/2 + +Cumulative product (`%Nx.Block.CumulativeProduct{}`). + +### Lowering + +Native `Torchx.cumulative_product/2`. + +## cumulative_min/2 + +Cumulative minimum (`%Nx.Block.CumulativeMin{}`). + +### Lowering + +Native `Torchx.cumulative_min/2`. + +## cumulative_max/2 + +Cumulative maximum (`%Nx.Block.CumulativeMax{}`). + +### Lowering + +Native `Torchx.cumulative_max/2`. + +## all_close/3 + +All-close comparison (`%Nx.Block.AllClose{}`). + +### Lowering + +Native `Torchx.all_close/4` after dtype merge of operands. + +### Options + +* `:rtol`, `:atol`, `:equal_nan` — passed to LibTorch + +## logical_not/1 + +Logical not (`%Nx.Block.LogicalNot{}`). + +### Lowering + +Native `Torchx.logical_not/1`. + +## phase/1 + +Complex phase (`%Nx.Block.Phase{}`). + +### Lowering + +Default Nx block callback (no dedicated LibTorch wrapper in `Torchx.Backend`). + +## runtime_call/4 + +Runtime Elixir callback from `defn`. + +### Behaviour + +Torchx does not provide a `defn` compiler. Outside `defn`, `runtime_call/4` +executes the callback directly with tensors on `Torchx.Backend`. Inside +`defn`, behaviour depends on the chosen compiler (`Nx.Defn.Evaluator` keeps +tensors on the active backend; EXLA has its own lowering). + +There is no LibTorch integration for device-side Elixir callbacks on Torchx. + +## backend_copy/2 + +Copy tensor data to another backend. + +### Behaviour + +Copying within `Torchx.Backend` uses `Torchx.to_device/2`. Copying to another +backend reads a binary blob via `Torchx.to_blob/1` and calls +`backend.from_binary/3`. + +## backend_transfer/2 + +Transfer tensor data to another backend. + +### Behaviour + +`backend_copy/2` followed by `backend_deallocate/1` on the source tensor. + +## backend_deallocate/1 + +Deallocate LibTorch tensor storage. + +### Behaviour + +Calls `Torchx.delete_tensor/1`. Returns `:already_deallocated` if the +underlying reference is invalid. + +## to_pointer/2 + +Export tensor memory as a pointer. + +### Behaviour + +**Not supported.** `to_pointer/2` raises on `Torchx.Backend`. + +## from_pointer/5 + +Import tensor memory from a pointer. + +### Behaviour + +**Not supported.** `from_pointer/5` raises on `Torchx.Backend`. + +## to_batched/3 + +Stream tensors in batches. + +### Behaviour + +Splits via `Torchx.split/2` on the leading axis. When `:leftover` is +`:repeat`, concatenates a tail slice so the final chunk matches `batch_size`. +May drop a partial last chunk when the size is not divisible and leftover is +not `:repeat`. diff --git a/torchx/guides/backend_documentation/nx_lin_alg.md b/torchx/guides/backend_documentation/nx_lin_alg.md new file mode 100644 index 0000000000..ee025b05ee --- /dev/null +++ b/torchx/guides/backend_documentation/nx_lin_alg.md @@ -0,0 +1,199 @@ +# Nx.LinAlg (Torchx) + +```elixir +iex> a = Nx.tensor([[4.0, 2.0], [2.0, 3.0]]) +iex> l = Nx.LinAlg.cholesky(a) +iex> Nx.shape(l) +{2, 2} +``` + +Torchx implementation notes for `Nx.LinAlg`. + +Each function below documents how Torchx handles the corresponding `Nx.LinAlg` +operation when `Torchx.Backend` is the active backend. + +Block-tagged operations are dispatched through `Torchx.Backend.block/4`. When +Torchx does not provide a native LibTorch path for a block, it invokes the +default `Nx.block/4` callback, which composes elementary Nx operations on the +same backend. + +Several linear algebra blocks are **not** supported natively on Apple MPS (`:mps`). +On that device, Torchx falls back to the default Nx block implementation instead +of calling LibTorch for LU, Eigh, Solve, Determinant, and Cholesky. `qr/2` and +`svd/2` copy inputs to CPU, run the native LibTorch kernel, and copy results +back to MPS. + +## cholesky/1 + +Cholesky decomposition (`%Nx.Block.LinAlg.Cholesky{}`). + +### Lowering + +On **CPU** and **CUDA**, Torchx calls `Torchx.cholesky/1` (LibTorch +`torch::linalg_cholesky`). + +On **MPS**, the default Nx callback is used because Cholesky is not in the +native MPS path. + +### Types + +Follows Nx promotion to floating point in the public API. LibTorch receives +the tensor's Torch scalar type (`:float`, `:double`, `:complex`, etc.). + +### Numerical notes + +Results may differ slightly from `Nx.BinaryBackend` due to LibTorch numerics +and rounding (half-to-even vs Elixir's half-away-from-zero elsewhere in Torchx). + +## solve/2 + +Linear system solve (`%Nx.Block.LinAlg.Solve{}`). + +### Lowering + +On **CPU** and **CUDA**, Torchx calls `Torchx.solve/2` (LibTorch +`torch::linalg_solve`). + +On **MPS**, the default Nx callback is used. + +### Singular matrices + +Before calling LibTorch, Torchx checks whether `|det(a)|` is below a dtype- +dependent epsilon (`1.0e-4` for `f32`, `1.0e-10` for `f64`). Singular systems +raise `ArgumentError` with `"can't solve for singular matrix"` instead of +letting LibTorch throw an opaque error. + +### Types + +Operands are cast to a merged floating-point Torch type before the solve. + +## qr/2 + +QR decomposition (`%Nx.Block.LinAlg.QR{}`). + +### Lowering + +Torchx calls `Torchx.qr/2`, mapping `mode: :reduced` to LibTorch's reduced QR +flag. + +On **MPS**, inputs are copied to CPU, QR is computed there, and `{q, r}` are +copied back to MPS. + +### Options + +* `:mode` — `:reduced` or `:complete`, passed to LibTorch +* `:eps` — ignored (used only by the default Nx callback on MPS fallback +paths for other blocks; QR always uses LibTorch when not on MPS) + +### Numerical notes + +`q` and `r` are not guaranteed to match other backends bit-for-bit. Verify +via reconstruction `q · r ≈ input`. + +## eigh/2 + +Hermitian eigendecomposition (`%Nx.Block.LinAlg.Eigh{}`). + +### Lowering + +On **CPU** and **CUDA**, integer inputs are promoted to `f32` before calling +`Torchx.eigh/1` (LibTorch `torch::linalg_eigh`). + +On **MPS**, the default Nx callback is used. + +### Options + +* `:max_iter` and `:eps` — honoured only by the default Nx callback (MPS +path); ignored by the LibTorch kernel + +### Numerical notes + +Eigenvectors are not unique. Prefer comparing reconstructions over direct +eigenvector equality. + +## svd/2 + +Singular value decomposition (`%Nx.Block.LinAlg.SVD{}`). + +### Lowering + +Torchx calls `Torchx.svd/2`, passing `full_matrices?` to LibTorch. + +On **MPS**, inputs are copied to CPU, SVD is computed there, and outputs are +copied back to MPS. + +Integer inputs are promoted to `f32` before the native call. + +### Types + +Complex SVD is **not** supported — `Nx.LinAlg.svd/2` raises for complex inputs +on Torchx. + +### Options + +* `:full_matrices?` — honoured by LibTorch +* `:max_iter` — ignored (LibTorch uses its own algorithm; the default Nx +iterative implementation is not used on the native path) + +### Numerical notes + +Singular values are the most stable quantity to compare across backends. + +## lu/1 + +LU decomposition with partial pivoting (`%Nx.Block.LinAlg.LU{}`). + +### Lowering + +On **CPU** and **CUDA**, integer inputs are promoted to `f32` before calling +`Torchx.lu/1` (LibTorch `torch::linalg_lu`). + +On **MPS**, the default Nx callback is used. + +### Types + +Permutation `p` keeps integer semantics from LibTorch; `l` and `u` are +floating point. + +## determinant/1 + +Matrix determinant (`%Nx.Block.LinAlg.Determinant{}`). + +### Lowering + +On **CPU** and **CUDA**, integer inputs are promoted to `f32` before calling +`Torchx.determinant/1`. + +On **MPS**, the default Nx callback is used. + +### Types + +Output dtype follows LibTorch / Nx promotion rules for the input. + +## triangular_solve/3 + +Triangular solve (direct `triangular_solve` callback, not an `Nx.block/4`). + +### Lowering + +Torchx calls `Torchx.triangular_solve/4` (LibTorch +`torch::linalg_solve_triangular`). + +On **MPS**, `a` and `b` are transferred to CPU for the solve; the result is +transferred back to MPS. + +### Options + +* `:lower` — honoured (`upper` flag inverted for LibTorch) +* `:transform_a` — `:none` and `:transpose` are supported +* `:left_side` — **only `true` is supported**; `left_side: false` raises +`ArgumentError` + +### Singular matrices + +Same determinant-based singularity check as `solve/2` before calling LibTorch. + +### Numerical notes + +Batched inputs are reshaped to `{batch, m, m}` and `{batch, m, n}` internally +when needed. diff --git a/torchx/lib/torchx.ex b/torchx/lib/torchx.ex index 506fa958ec..e94410e6a3 100644 --- a/torchx/lib/torchx.ex +++ b/torchx/lib/torchx.ex @@ -157,6 +157,14 @@ defmodule Torchx do PyTorch implements a variety of devices, which can be seen below. #{@valid_devices_md_list} + + ## Backend documentation + + Per-operation implementation notes for Nx linear algebra are documented in + [Nx.LinAlg](backend_documentation-nx_lin_alg.html). Top-level `Nx` blocks and + transfers are documented in [Nx](backend_documentation-nx.html). See the + [backend documentation convention](https://hexdocs.pm/nx/backend_documentation-convention.html) + in the Nx package. """ use Torchx.Macro alias Torchx.NIF diff --git a/torchx/lib/torchx/backend.ex b/torchx/lib/torchx/backend.ex index 3bbc195b0e..53324dc767 100644 --- a/torchx/lib/torchx/backend.ex +++ b/torchx/lib/torchx/backend.ex @@ -27,6 +27,11 @@ defmodule Torchx.Backend do and nan becomes zero. `Torchx` behaviour is type dependent with no clear rule across types. + Per-operation implementation notes are documented in the + [Backend documentation](backend_documentation.html) guides — see the + [backend documentation convention](https://hexdocs.pm/nx/backend_documentation-convention.html) + in the Nx package. + ## Options * `:device` - Defaults to `Torchx.default_device/0`. An atom representing the diff --git a/torchx/mix.exs b/torchx/mix.exs index d53c27bc34..43ecd656b3 100644 --- a/torchx/mix.exs +++ b/torchx/mix.exs @@ -58,7 +58,22 @@ defmodule Torchx.MixProject do main: "Torchx", source_url_pattern: "#{@source_url}/blob/v#{@version}/torchx/%{path}#L%{line}", extras: [ - "CHANGELOG.md" + "CHANGELOG.md", + "guides/backend_documentation/index.md": [ + filename: "backend_documentation", + title: "Backend documentation" + ], + "guides/backend_documentation/nx.md": [ + filename: "backend_documentation-nx", + title: "Nx (Torchx)" + ], + "guides/backend_documentation/nx_lin_alg.md": [ + filename: "backend_documentation-nx_lin_alg", + title: "Nx.LinAlg (Torchx)" + ] + ], + groups_for_extras: [ + "Backend documentation": ~r"^guides/backend_documentation/" ] ] end @@ -68,7 +83,7 @@ defmodule Torchx.MixProject do maintainers: ["Paulo Valente", "José Valim"], licenses: ["Apache-2.0"], links: %{"GitHub" => @source_url}, - files: ["lib", "mix.exs", "README.md", "LICENSE", "CHANGELOG.md", "c_src", "CMakeLists.txt"] + files: ["lib", "mix.exs", "README.md", "LICENSE", "CHANGELOG.md", "guides", "c_src", "CMakeLists.txt"] ] end diff --git a/torchx/test/torchx/backend_documentation_test.exs b/torchx/test/torchx/backend_documentation_test.exs new file mode 100644 index 0000000000..8347e9d18e --- /dev/null +++ b/torchx/test/torchx/backend_documentation_test.exs @@ -0,0 +1,7 @@ +defmodule Torchx.BackendDocumentationTest do + use ExUnit.Case, async: true + + doctest_file "guides/backend_documentation/index.md" + doctest_file "guides/backend_documentation/nx.md" + doctest_file "guides/backend_documentation/nx_lin_alg.md" +end From c04e09f410c3ec2ef58e54e5569ca56e0e91cbe1 Mon Sep 17 00:00:00 2001 From: Chapaman <14204271+Chapaman@users.noreply.github.com> Date: Wed, 1 Jul 2026 22:26:28 -0300 Subject: [PATCH 2/6] fix formatting --- exla/test/exla/backend_documentation_test.exs | 6 +++--- nx/test/nx/backend_documentation_test.exs | 2 +- torchx/mix.exs | 11 ++++++++++- torchx/test/torchx/backend_documentation_test.exs | 6 +++--- 4 files changed, 17 insertions(+), 8 deletions(-) diff --git a/exla/test/exla/backend_documentation_test.exs b/exla/test/exla/backend_documentation_test.exs index b6de7ca3cd..081319db10 100644 --- a/exla/test/exla/backend_documentation_test.exs +++ b/exla/test/exla/backend_documentation_test.exs @@ -1,7 +1,7 @@ defmodule EXLA.BackendDocumentationTest do use ExUnit.Case, async: true - doctest_file "guides/backend_documentation/index.md" - doctest_file "guides/backend_documentation/nx.md" - doctest_file "guides/backend_documentation/nx_lin_alg.md" + doctest_file("guides/backend_documentation/index.md") + doctest_file("guides/backend_documentation/nx.md") + doctest_file("guides/backend_documentation/nx_lin_alg.md") end diff --git a/nx/test/nx/backend_documentation_test.exs b/nx/test/nx/backend_documentation_test.exs index 3dfc6d4ae4..f90119a44d 100644 --- a/nx/test/nx/backend_documentation_test.exs +++ b/nx/test/nx/backend_documentation_test.exs @@ -1,5 +1,5 @@ defmodule Nx.BackendDocumentationTest do use ExUnit.Case, async: true - doctest_file "guides/backend_documentation/convention.md" + doctest_file("guides/backend_documentation/convention.md") end diff --git a/torchx/mix.exs b/torchx/mix.exs index 43ecd656b3..6e842be44f 100644 --- a/torchx/mix.exs +++ b/torchx/mix.exs @@ -83,7 +83,16 @@ defmodule Torchx.MixProject do maintainers: ["Paulo Valente", "José Valim"], licenses: ["Apache-2.0"], links: %{"GitHub" => @source_url}, - files: ["lib", "mix.exs", "README.md", "LICENSE", "CHANGELOG.md", "guides", "c_src", "CMakeLists.txt"] + files: [ + "lib", + "mix.exs", + "README.md", + "LICENSE", + "CHANGELOG.md", + "guides", + "c_src", + "CMakeLists.txt" + ] ] end diff --git a/torchx/test/torchx/backend_documentation_test.exs b/torchx/test/torchx/backend_documentation_test.exs index 8347e9d18e..ed0140fac5 100644 --- a/torchx/test/torchx/backend_documentation_test.exs +++ b/torchx/test/torchx/backend_documentation_test.exs @@ -1,7 +1,7 @@ defmodule Torchx.BackendDocumentationTest do use ExUnit.Case, async: true - doctest_file "guides/backend_documentation/index.md" - doctest_file "guides/backend_documentation/nx.md" - doctest_file "guides/backend_documentation/nx_lin_alg.md" + doctest_file("guides/backend_documentation/index.md") + doctest_file("guides/backend_documentation/nx.md") + doctest_file("guides/backend_documentation/nx_lin_alg.md") end From e21b41f24a14fd0de48ebb4d2a276b4cbef8c8ff Mon Sep 17 00:00:00 2001 From: Chapaman <14204271+Chapaman@users.noreply.github.com> Date: Mon, 6 Jul 2026 17:28:32 -0300 Subject: [PATCH 3/6] implement josevalim and paulovalente feedback --- exla/guides/backend_documentation/index.md | 8 +- exla/guides/backend_documentation/nx.md | 180 +----------------- .../backend_documentation/nx_lin_alg.md | 164 ++-------------- exla/lib/exla.ex | 6 +- exla/mix.exs | 9 +- nx/guides/backend_documentation/convention.md | 35 ++-- nx/lib/nx.ex | 70 +------ nx/lib/nx/backend.ex | 5 +- nx/lib/nx/lin_alg.ex | 15 +- nx/mix.exs | 3 +- torchx/guides/backend_documentation/index.md | 8 +- torchx/guides/backend_documentation/nx.md | 147 +------------- .../backend_documentation/nx_lin_alg.md | 78 ++------ torchx/lib/torchx.ex | 7 +- torchx/lib/torchx/backend.ex | 2 +- torchx/mix.exs | 9 +- 16 files changed, 88 insertions(+), 658 deletions(-) diff --git a/exla/guides/backend_documentation/index.md b/exla/guides/backend_documentation/index.md index 32fb157aa3..bd978f1b10 100644 --- a/exla/guides/backend_documentation/index.md +++ b/exla/guides/backend_documentation/index.md @@ -2,14 +2,14 @@ EXLA-specific documentation for Nx backend behaviour. -These guides describe how EXLA lowers and executes Nx operations. They mirror -the structure of the Nx API (see the [backend documentation convention](https://hexdocs.pm/nx/backend_documentation-convention.html)) -but are **not callable** — use `Nx` and `compiler: EXLA` in your code. +These guides document divergent behaviour, backend-specific options, and +limitations for Nx operations. They mirror the structure of the Nx API — see +the [backend documentation convention](https://hexdocs.pm/nx/backend_documentation-convention.html). ## Guides * [Nx](backend_documentation-nx.html) — top-level `Nx` blocks, transfers, and `defn` integration - * [Nx.LinAlg](backend_documentation-nx_lin_alg.html) — linear algebra blocks and related lowerings + * [Nx.LinAlg](backend_documentation-nx_lin_alg.html) — linear algebra blocks Implementation code lives in `EXLA.Defn`, `EXLA.CustomCall`, and related modules. diff --git a/exla/guides/backend_documentation/nx.md b/exla/guides/backend_documentation/nx.md index d510b06945..a1e9d1e00f 100644 --- a/exla/guides/backend_documentation/nx.md +++ b/exla/guides/backend_documentation/nx.md @@ -8,174 +8,19 @@ iex> Nx.take(Nx.tensor([10, 20, 30]), Nx.tensor([0, 2])) > ``` -EXLA implementation notes for top-level `Nx` operations. - -Each function below documents how EXLA handles the corresponding `Nx` API when -`compiler: EXLA` is used (or when `EXLA.Backend` stores tensors). +EXLA-specific notes for top-level `Nx` operations where behaviour diverges from +the portable Nx API or from other backends. Linear algebra is documented separately in [Nx.LinAlg](backend_documentation-nx_lin_alg.html). -Many `Nx` functions delegate to `Nx.block/4`. During EXLA compilation, -specialized lowerings exist for some block tags; otherwise EXLA compiles the -default callback as an XLA subcomputation. - -## take/3 - -Tensor indexing (`%Nx.Block.Take{}`). - -### Lowering - -Lowered to StableHLO `gather` in `EXLA.Defn` (not a custom call). - -### Options - -* `:axis` — honoured; indices are gathered along this axis - -### Platforms - -All EXLA clients. - -## take_along_axis/3 - -Take along axis (`%Nx.Block.TakeAlongAxis{}`). - -### Lowering - -No dedicated gather lowering — EXLA compiles the default Nx callback as an XLA -subcomputation. - -### Options - -* `:axis` — honoured by the default callback - -## top_k/2 - -Top-k values and indices (`%Nx.Block.TopK{}`). - -### Lowering - -Lowered to StableHLO `top_k` via `EXLA.MLIR.Value.top_k/3`. - -### Options - -* `:k` — number of elements to return per slice - -## fft2/2 - -Two-dimensional FFT (`%Nx.Block.FFT2{}`). - -### Lowering - -Lowered to StableHLO FFT ops via `EXLA.MLIR.Value.fft/4` with mode `:fft`. - -### Options - -* `:lengths`, `:axes`, `:eps` — forwarded to the lowering; `:eps` may be used -for cleanup in the default callback when compilation falls back - -## ifft2/2 - -Two-dimensional inverse FFT (`%Nx.Block.IFFT2{}`). - -### Lowering - -Lowered to StableHLO FFT ops with mode `:ifft`. - -## rfft/2 - -Real FFT (`%Nx.Block.RFFT{}`). - -### Lowering - -Lowered to StableHLO real FFT via `Value.fft/4` with mode `:rfft`. Input is -treated as real; output type is complex per the expression template. - -## irfft/2 - -Inverse real FFT (`%Nx.Block.IRFFT{}`). - -### Lowering - -Lowered to StableHLO IRFFT via `Value.fft/4` with mode `:irfft`. - -## cumulative_sum/2 - -Cumulative sum (`%Nx.Block.CumulativeSum{}`). - -### Lowering - -Default callback compiled as an XLA subcomputation (no native custom call). - -### Options - -* `:axis`, `:reverse` — honoured by the default callback - -## cumulative_product/2 - -Cumulative product (`%Nx.Block.CumulativeProduct{}`). - -### Lowering - -Default callback compiled as an XLA subcomputation. - -## cumulative_min/2 - -Cumulative minimum (`%Nx.Block.CumulativeMin{}`). - -### Lowering - -Default callback compiled as an XLA subcomputation. - -## cumulative_max/2 - -Cumulative maximum (`%Nx.Block.CumulativeMax{}`). - -### Lowering - -Default callback compiled as an XLA subcomputation. - -## all_close/3 - -All-close comparison (`%Nx.Block.AllClose{}`). - -### Lowering - -Default callback compiled as an XLA subcomputation. - -### Options - -* `:rtol`, `:atol`, `:equal_nan` — honoured by the default callback - -## logical_not/1 - -Logical not (`%Nx.Block.LogicalNot{}`). - -### Lowering - -Default callback compiled as an XLA subcomputation (delegates to element-wise -equality in the reference implementation). - -## phase/1 - -Complex phase (`%Nx.Block.Phase{}`). - -### Lowering - -Default callback compiled as an XLA subcomputation. - ## runtime_call/4 Runtime Elixir callback from `defn` (`runtime_call` expression). -### Lowering - -On **host** and **CUDA** clients, EXLA emits a device-side callback bridged -through `EXLA.MLIR.Value.runtime_call/3`. Callback tensors are materialized for -the Elixir function; results must match the output template. - ### Platforms -* **Host / CUDA** — supported +* **Host / CUDA** — supported via a device-side callback bridged through +`EXLA.MLIR.Value.runtime_call/3` * **ROCm / TPU** — raises with a message to use `:host` or `:cuda` ### Outside `defn` @@ -212,13 +57,9 @@ Transfer tensor data to another backend (`EXLA.Backend.backend_transfer/3`). Same as `backend_copy/2` followed by deallocation of the source device buffer when leaving `EXLA.Backend`. -## backend_deallocate/1 - -Deallocate device memory (`EXLA.Backend.backend_deallocate/1`). - -### Behaviour +### Options -Calls `EXLA.DeviceBuffer.deallocate/1` for tensors on `EXLA.Backend`. +* `:client`, `:device_id` — select the EXLA client and device ## to_pointer/2 @@ -254,11 +95,6 @@ or **CUDA** clients. Pointer `data_size` must match the tensor byte size. * `:client`, `:device_id` — destination device * `:names` — tensor names for the result template -## to_batched/3 - -Stream tensors in batches (`EXLA.Backend.to_batched/3`). - -### Behaviour +### Limitations -Splits the leading axis via XLA slice and concatenate operations on device -buffers. Supports `:leftover` `:repeat` and `:discard` like `Nx.to_batched/2`. +Not supported for ROCm or TPU. diff --git a/exla/guides/backend_documentation/nx_lin_alg.md b/exla/guides/backend_documentation/nx_lin_alg.md index e61a16e5fb..7c00d62cdd 100644 --- a/exla/guides/backend_documentation/nx_lin_alg.md +++ b/exla/guides/backend_documentation/nx_lin_alg.md @@ -7,123 +7,46 @@ iex> Nx.shape(l) {2, 2} ``` -EXLA implementation notes for `Nx.LinAlg`. - -Each function below documents how EXLA handles the corresponding `Nx.LinAlg` -operation when `compiler: EXLA` is used. - -Operations that use `Nx.block/4` are dispatched through `EXLA.Defn`. When -`EXLA.CustomCall.call/4` returns `:skip`, EXLA compiles the default block -callback as an XLA subcomputation. When it returns `{:ok, spec}`, EXLA emits a -StableHLO `custom_call` instead. +EXLA-specific notes for `Nx.LinAlg` where behaviour diverges from the portable +Nx API or from other backends. On GPU and TPU clients, the `:precision` option passed to `Nx.Defn.jit/2` affects accumulator precision for many composed linear algebra graphs compiled from the default callbacks. For decomposition verification tests, prefer `precision: :highest`. -## cholesky/1 - -Cholesky decomposition (`%Nx.Block.LinAlg.Cholesky{}`). - -### Lowering - -EXLA does **not** provide a native custom call for Cholesky. The default -callback (`Nx.LinAlg.Cholesky.cholesky/1`) is compiled as an XLA -subcomputation. - -### Supported platforms - -All EXLA clients (`:host`, `:cuda`, `:rocm`, `:tpu`) use the compiled default -implementation. - -### Types - -Follows Nx promotion rules: integer and unsigned inputs are promoted to -floating point in the public API before the block is invoked. - -### Numerical notes - -Positive-definiteness is not validated at compile time. Ill-conditioned inputs -may produce `NaN` values without raising, consistent with the reference Nx -implementation. - -## solve/2 - -Linear system solve (`%Nx.Block.LinAlg.Solve{}`). - -### Lowering - -EXLA compiles the default callback, which expresses `solve/2` as an LU -factorization followed by triangular solves. - -### Supported platforms - -All EXLA clients use the compiled default implementation. - -### Types - -Output type is floating point per `Nx.LinAlg.solve/2`. The solve block depends -on `lu/1` and `triangular_solve/3` inside the compiled graph. - ## qr/2 QR decomposition (`%Nx.Block.LinAlg.QR{}`). -### Lowering - -On the **host** client, with **real** input element types, EXLA emits a -native CPU custom call via `EXLA.CustomCall` when the dtype is supported: - -* `{:f, 32}` → `qr_cpu_custom_call_f32` -* `{:f, 64}` → `qr_cpu_custom_call_f64` -* `{:f, 16}` → `qr_cpu_custom_call_f16` -* `{:bf, 16}` → `qr_cpu_custom_call_bf16` -* `{:s, _}` and `{:u, _}` → `qr_cpu_custom_call_f32` with operand cast to -`f32` - -Otherwise EXLA compiles the default callback (`Nx.LinAlg.QR.qr/2`). - ### Supported platforms -* **Host** — native custom call for supported real dtypes; default callback -for complex inputs and unsupported types +* **Host** — native CPU custom call for supported **real** dtypes (`f16`, +`bf16`, `f32`, `f64`, and integer types cast to `f32`); default callback for +complex inputs and unsupported types * **CUDA / ROCm / TPU** — default callback only (no native QR custom call) ### Options -* `:mode` — honoured by whichever lowering is selected (`:reduced` or +* `:mode` — honoured by whichever implementation is selected (`:reduced` or `:complete`) -* `:eps` — used by the default callback only; ignored by the native custom -call +* `:eps` — used by the default callback only; ignored by the native custom call ### Numerical notes -Native and default lowerings may differ slightly in orthogonality of `q` and -triangular structure of `r` within floating-point tolerance. Reconstruction +Native and default implementations may differ slightly in orthogonality of `q` +and triangular structure of `r` within floating-point tolerance. Reconstruction `q · r` should match the input within reasonable tolerances for the dtype. ## eigh/2 Hermitian eigendecomposition (`%Nx.Block.LinAlg.Eigh{}`). -### Lowering - -On the **host** client, with **real** input element types, EXLA emits a -native CPU custom call when the dtype is supported: - -* `{:f, 32}` → `eigh_cpu_custom_call_f32` -* `{:f, 64}` → `eigh_cpu_custom_call_f64` -* `{:s, _}` and `{:u, _}` → `eigh_cpu_custom_call_f32` with operand cast to -`f32` - -Otherwise EXLA compiles the default callback (`Nx.LinAlg.BlockEigh.eigh/2` or -`Nx.LinAlg.Eigh.eigh/2` depending on input rank). - ### Supported platforms -* **Host** — native custom call for supported real dtypes; default callback -for complex inputs and unsupported types +* **Host** — native CPU custom call for supported **real** dtypes (`f32`, +`f64`, and integer types cast to `f32`); default callback for complex inputs and +unsupported types * **CUDA / ROCm / TPU** — default callback only ### Options @@ -141,15 +64,6 @@ across backends. Singular value decomposition (`%Nx.Block.LinAlg.SVD{}`). -### Lowering - -EXLA compiles the default callback (`Nx.LinAlg.SVD.svd/2`). There is no native -SVD custom call. - -### Supported platforms - -All EXLA clients use the compiled default implementation. - ### Options * `:max_iter` — honoured (default iterative algorithm in Nx) @@ -160,57 +74,3 @@ All EXLA clients use the compiled default implementation. SVD is iterative; results may vary across devices and precision settings. Singular values are the most stable quantity to compare across backends. Use `precision: :highest` on GPU when verifying decompositions. - -## lu/1 - -LU decomposition with partial pivoting (`%Nx.Block.LinAlg.LU{}`). - -### Lowering - -EXLA compiles the default callback (`Nx.LinAlg.LU.lu/1`) as an XLA -subcomputation. Permutation matrix `p` uses integer (`s32`) type; `l` and `u` -are floating point. - -### Supported platforms - -All EXLA clients use the compiled default implementation. - -### Numerical notes - -For rank-deficient inputs, `u` may contain zeros on the diagonal. This matches -the reference Nx behaviour. - -## determinant/1 - -Matrix determinant (`%Nx.Block.LinAlg.Determinant{}`). - -### Lowering - -EXLA compiles the default callback. Small matrices use closed-form expressions; -larger matrices use LU-based determinant computation inside the compiled graph. - -### Supported platforms - -All EXLA clients use the compiled default implementation. - -### Types - -Output is floating point. Complex determinants are supported when the input is -complex. - -## triangular_solve/3 - -Triangular solve (direct `triangular_solve` callback, not an `Nx.block/4`). - -### Lowering - -Lowered to the StableHLO `triangular_solve` op via `EXLA.MLIR.Value.triangular_solve/4`. - -### Supported platforms - -All EXLA clients. - -### Options - -EXLA honours `lower:`, `left_side:`, and `transform_a:` as defined in -`Nx.LinAlg.triangular_solve/3`. diff --git a/exla/lib/exla.ex b/exla/lib/exla.ex index cb8d162e50..5b019a57d9 100644 --- a/exla/lib/exla.ex +++ b/exla/lib/exla.ex @@ -82,9 +82,9 @@ defmodule EXLA do your block tag struct; see `EXLA.CustomCall` for the `call/4` contract, including returning `:skip` to fall back to the block's default Elixir callback. - Per-operation lowering notes for Nx functions are documented in the - [Backend documentation](backend_documentation.html) guides (for example - [Nx](backend_documentation-nx.html) and + Backend-specific behaviour, options, and limitations for Nx functions are + documented in the [Backend documentation](backend_documentation.html) guides + (for example [Nx](backend_documentation-nx.html) and [Nx.LinAlg](backend_documentation-nx_lin_alg.html)). See also the [backend documentation convention](https://hexdocs.pm/nx/backend_documentation-convention.html) in the Nx package. diff --git a/exla/mix.exs b/exla/mix.exs index 9dc28f4cbe..59c06b0473 100644 --- a/exla/mix.exs +++ b/exla/mix.exs @@ -90,16 +90,13 @@ defmodule EXLA.MixProject do "guides/rotating-image.livemd", "CHANGELOG.md", "guides/backend_documentation/index.md": [ - filename: "backend_documentation", - title: "Backend documentation" + filename: "backend_documentation" ], "guides/backend_documentation/nx.md": [ - filename: "backend_documentation-nx", - title: "Nx (EXLA)" + filename: "backend_documentation-nx" ], "guides/backend_documentation/nx_lin_alg.md": [ - filename: "backend_documentation-nx_lin_alg", - title: "Nx.LinAlg (EXLA)" + filename: "backend_documentation-nx_lin_alg" ] ], skip_undefined_reference_warnings_on: ["CHANGELOG.md"], diff --git a/nx/guides/backend_documentation/convention.md b/nx/guides/backend_documentation/convention.md index c9df6ffea3..d21cabe27a 100644 --- a/nx/guides/backend_documentation/convention.md +++ b/nx/guides/backend_documentation/convention.md @@ -2,14 +2,13 @@ Nx exposes a stable public API (`Nx`, `Nx.LinAlg`, and so on) while individual backends implement the underlying callbacks and optional `Nx.block/4` lowerings. -Users who need to know *how* a specific backend executes an operation should +Users who need to know how a specific backend differs from the portable API should not have to read backend source code. ## Documentation guides Each backend may publish **backend documentation guides** that mirror the Nx -module structure and document backend-specific behaviour for each callback. -The naming convention is: +module structure. The naming convention is: guides/backend_documentation/nx.md # mirrors top-level `Nx` guides/backend_documentation/nx_lin_alg.md # mirrors `Nx.LinAlg` @@ -18,30 +17,32 @@ For example, EXLA documents operations in its [Backend documentation](https://he guides, and Torchx in its [Backend documentation](https://hexdocs.pm/torchx/backend_documentation.html) guides. -These guides describe primitives that backends implement, without being part of the -callable Nx API. - ## What to document -Backends should document, for each mirrored function: +Document an operation only when a backend has **divergent behaviour**, **backend-specific +options**, or **trade-offs and limitations** worth calling out relative to the +portable Nx API or to other backends. Examples: + + * an option that a backend honours, ignores, or overrides + * platform or device support gaps (for example an operation unavailable on MPS) + * numerical or performance trade-offs that affect how you use the API + * warnings about deadlock, unsupported pointer modes, and similar constraints - * whether the operation is lowered natively, compiled from the default - `Nx.block/4` callback, or handled through another mechanism (for example - a direct StableHLO op) - * supported types, devices, and platforms - * options that the backend honours, ignores, or overrides - * numerical behaviour, performance trade-offs, and known limitations relative - to the reference implementation in Nx +Do not document routine lowerings (for example which StableHLO op or LibTorch +kernel is used) when behaviour matches the reference Nx implementation. If you +are implementing or overriding custom blocks, the backend source remains the best +source of truth. -Cross-link from the Nx function doc when appropriate (for example -`See EXLA and Torchx backend documentation for qr/2`). +Cross-link from the Nx function doc when a backend guide has relevant notes for +that operation. ## Relation to `Nx.Backend` Most tensor operations go through `Nx.Backend` callbacks. Operations that need a portable default with optional native acceleration use `Nx.block/4` and the `Nx.Block.*` tags — see `Nx.Backend.block/4`. Backend documentation guides -should cover both callback implementations and block lowerings. +should cover divergent behaviour for both callback implementations and block +lowerings. ## Example diff --git a/nx/lib/nx.ex b/nx/lib/nx.ex index f6a8fb6c10..a70fb73d45 100644 --- a/nx/lib/nx.ex +++ b/nx/lib/nx.ex @@ -41,7 +41,8 @@ defmodule Nx do Backend-specific notes for extensible operations (`Nx.block/4`, transfers, and related APIs) are documented in the [backend documentation convention](backend_documentation-convention.html) - guide and in each backend's **Backend documentation** pages on HexDocs (for example + guide and in each backend's **Backend documentation** pages on HexDocs when a + backend has divergent behaviour or backend-specific options (for example [EXLA](https://hexdocs.pm/exla/backend_documentation.html) and [Torchx](https://hexdocs.pm/torchx/backend_documentation.html)). @@ -1622,11 +1623,6 @@ defmodule Nx do iex> Nx.take_diagonal(Nx.iota({3, 3}), offset: -4) ** (ArgumentError) absolute value of offset must be less than length of axis 0 when negative, got: -4 - - ## Backend notes - - See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#take-3) and - [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#take-3). """ @doc type: :creation def take_diagonal(tensor, opts \\ []) do @@ -2533,10 +2529,6 @@ defmodule Nx do [3, 1] > - ## Backend notes - - See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#to_batched-3) and - [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#to_batched-3). """ @doc type: :conversion def to_batched(tensor, batch_size, opts \\ []) @@ -5082,10 +5074,6 @@ defmodule Nx do Note: This function cannot be used in `defn`. - ## Backend notes - - See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#backend_deallocate-1) and - [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#backend_deallocate-1). """ @doc type: :backend def backend_deallocate(tensor_or_container) do @@ -7261,10 +7249,6 @@ defmodule Nx do > - ## Backend notes - - See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#logical_not-1) and - [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#logical_not-1). """ @doc type: :element def logical_not(tensor) do @@ -8735,8 +8719,7 @@ defmodule Nx do ## Backend notes - See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#phase-1) and - [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#phase-1). + See [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#phase-1). """ @doc type: :element def phase(tensor) do @@ -9324,10 +9307,6 @@ defmodule Nx do ] > - ## Backend notes - - See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#all_close-3) and - [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#all_close-3). """ @doc type: :aggregation def all_close(a, b, opts \\ []) do @@ -11733,10 +11712,6 @@ defmodule Nx do ] > - ## Backend notes - - See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#cumulative_sum-2) and - [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#cumulative_sum-2). """ @doc type: :cumulative def cumulative_sum(tensor, opts \\ []), @@ -11814,10 +11789,6 @@ defmodule Nx do ] > - ## Backend notes - - See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#cumulative_product-2) and - [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#cumulative_product-2). """ @doc type: :cumulative def cumulative_product(tensor, opts \\ []), @@ -11895,10 +11866,6 @@ defmodule Nx do ] > - ## Backend notes - - See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#cumulative_min-2) and - [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#cumulative_min-2). """ @doc type: :cumulative def cumulative_min(tensor, opts \\ []), @@ -11976,10 +11943,6 @@ defmodule Nx do ] > - ## Backend notes - - See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#cumulative_max-2) and - [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#cumulative_max-2). """ @doc type: :cumulative def cumulative_max(tensor, opts \\ []), @@ -14912,11 +14875,6 @@ defmodule Nx do iex> idx = Nx.tensor([[2.0], [1.0], [2.0]], type: :f32) iex> Nx.take_along_axis(tensor, idx, axis: 1) ** (ArgumentError) indices must be an integer tensor, got {:f, 32} - - ## Backend notes - - See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#take_along_axis-3) and - [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#take_along_axis-3). """ @doc type: :indexed def take_along_axis(tensor, indices, opts \\ []) when is_list(opts) do @@ -15745,10 +15703,6 @@ defmodule Nx do ** (ArgumentError) top_k input must have at least rank 1 - ## Backend notes - - See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#top_k-2) and - [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#top_k-2). """ @doc type: :ndim def top_k(tensor, opts \\ []) do @@ -16987,8 +16941,7 @@ defmodule Nx do ## Backend notes - See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#fft2-2) and - [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#fft2-2). + See [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#fft2-2). """ @doc type: :ndim def fft2(tensor, opts \\ []), do: call_fft2(tensor, opts, :fft2) @@ -17147,11 +17100,6 @@ defmodule Nx do iex> Nx.ifft2(Nx.tensor([1, 1]), length: :invalid) ** (ArgumentError) expected a tensor with rank > 1, got tensor with rank 1 - - ## Backend notes - - See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#ifft2-2) and - [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#ifft2-2). """ @doc type: :ndim def ifft2(tensor, opts \\ []), do: call_fft2(tensor, opts, :ifft2) @@ -17326,11 +17274,6 @@ defmodule Nx do iex> Nx.rfft(Nx.tensor([1.0, 1.0]), length: :invalid) ** (ArgumentError) expected an integer or :power_of_two as length, got: :invalid - - ## Backend notes - - See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#rfft-2) and - [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#rfft-2). """ @doc type: :ndim def rfft(tensor, opts \\ []) do @@ -17452,11 +17395,6 @@ defmodule Nx do iex> Nx.irfft(Nx.tensor([1.0, 1.0]), length: :invalid) ** (ArgumentError) expected a positive integer as length, got: :invalid - - ## Backend notes - - See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#irfft-2) and - [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#irfft-2). """ @doc type: :ndim def irfft(tensor, opts \\ []) do diff --git a/nx/lib/nx/backend.ex b/nx/lib/nx/backend.ex index bc91bbdd02..7d8c938786 100644 --- a/nx/lib/nx/backend.ex +++ b/nx/lib/nx/backend.ex @@ -30,8 +30,9 @@ defmodule Nx.Backend do across backends. Backends may publish backend documentation guides that mirror the Nx API and - describe implementation specifics — see the - [backend documentation convention](backend_documentation-convention.html) guide. + describe divergent behaviour, backend-specific options, and limitations — see + the [backend documentation convention](backend_documentation-convention.html) + guide. """ @type t :: %{__struct__: atom()} diff --git a/nx/lib/nx/lin_alg.ex b/nx/lib/nx/lin_alg.ex index ee9e7c00d3..e115b60dd5 100644 --- a/nx/lib/nx/lin_alg.ex +++ b/nx/lib/nx/lin_alg.ex @@ -139,8 +139,7 @@ defmodule Nx.LinAlg do ## Backend notes - See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx_lin_alg.html#cholesky-1) and - [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx_lin_alg.html#cholesky-1). + See [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx_lin_alg.html#cholesky-1). """ def cholesky(tensor) do %T{vectorized_axes: vectorized_axes} = tensor = Nx.to_tensor(tensor) @@ -592,8 +591,7 @@ defmodule Nx.LinAlg do ## Backend notes - See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx_lin_alg.html#triangular_solve-3) and - [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx_lin_alg.html#triangular_solve-3). + See [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx_lin_alg.html#triangular_solve-3). """ def triangular_solve(a, b, opts \\ []) do opts = keyword!(opts, lower: true, left_side: true, transform_a: :none) @@ -710,8 +708,7 @@ defmodule Nx.LinAlg do ## Backend notes - See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx_lin_alg.html#solve-2) and - [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx_lin_alg.html#solve-2). + See [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx_lin_alg.html#solve-2). """ # IMPORTANT: This function cannot be a defn because # optional needs to work on the actual backend. @@ -1769,8 +1766,7 @@ defmodule Nx.LinAlg do ## Backend notes - See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx_lin_alg.html#lu-1) and - [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx_lin_alg.html#lu-1). + See [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx_lin_alg.html#lu-1). """ def lu(tensor) do %T{vectorized_axes: vectorized_axes} = tensor = Nx.to_tensor(tensor) @@ -2023,8 +2019,7 @@ defmodule Nx.LinAlg do ## Backend notes - See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx_lin_alg.html#determinant-1) and - [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx_lin_alg.html#determinant-1). + See [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx_lin_alg.html#determinant-1). """ # IMPORTANT: This function cannot be a defn because # optional needs to work on the actual backend. diff --git a/nx/mix.exs b/nx/mix.exs index df9585e892..f6c3c24e22 100644 --- a/nx/mix.exs +++ b/nx/mix.exs @@ -76,8 +76,7 @@ defmodule Nx.MixProject do "guides/advanced/complex_fft.livemd", "guides/exercises/exercises-1-20.livemd", "guides/backend_documentation/convention.md": [ - filename: "backend_documentation-convention", - title: "Backend documentation convention" + filename: "backend_documentation-convention" ] ], skip_undefined_reference_warnings_on: ["CHANGELOG.md"], diff --git a/torchx/guides/backend_documentation/index.md b/torchx/guides/backend_documentation/index.md index c056db57b0..80a8a7c36d 100644 --- a/torchx/guides/backend_documentation/index.md +++ b/torchx/guides/backend_documentation/index.md @@ -2,14 +2,14 @@ Torchx-specific documentation for Nx backend behaviour. -These guides describe how Torchx executes Nx operations through LibTorch. -They mirror the structure of the Nx API (see the [backend documentation convention](https://hexdocs.pm/nx/backend_documentation-convention.html)) -but are **not callable** — use `Nx` with `Torchx.Backend` in your code. +These guides document divergent behaviour, backend-specific options, and +limitations for Nx operations. They mirror the structure of the Nx API — see +the [backend documentation convention](https://hexdocs.pm/nx/backend_documentation-convention.html). ## Guides * [Nx](backend_documentation-nx.html) — top-level `Nx` blocks, transfers, and callbacks - * [Nx.LinAlg](backend_documentation-nx_lin_alg.html) — linear algebra blocks and related callbacks + * [Nx.LinAlg](backend_documentation-nx_lin_alg.html) — linear algebra blocks Implementation code lives in `Torchx.Backend` and the `Torchx` NIF bindings. diff --git a/torchx/guides/backend_documentation/nx.md b/torchx/guides/backend_documentation/nx.md index 9998d71b90..448136b4b3 100644 --- a/torchx/guides/backend_documentation/nx.md +++ b/torchx/guides/backend_documentation/nx.md @@ -8,10 +8,8 @@ iex> Nx.take(Nx.tensor([10, 20, 30]), Nx.tensor([0, 2])) > ``` -Torchx implementation notes for top-level `Nx` operations. - -Each function below documents how Torchx handles the corresponding `Nx` API -when `Torchx.Backend` is active. +Torchx-specific notes for top-level `Nx` operations where behaviour diverges from +the portable Nx API or from other backends. Linear algebra is documented separately in [Nx.LinAlg](backend_documentation-nx_lin_alg.html). @@ -19,139 +17,22 @@ Torchx is an **Nx backend**, not a `defn` compiler. Block lowerings run through `Torchx.Backend.block/4` during eager execution or when another compiler (such as `Nx.Defn.Evaluator`) invokes backend callbacks. -## take/3 - -Tensor indexing (`%Nx.Block.Take{}`). - -### Lowering - -No native LibTorch path — Torchx invokes the default Nx block callback, which -composes `Nx` gather operations on the Torchx backend. - -### Options - -* `:axis` — honoured by the default callback - -## take_along_axis/3 - -Take along axis (`%Nx.Block.TakeAlongAxis{}`). - -### Lowering - -Native `Torchx.gather/2` along the configured axis (LibTorch `gather`). - -## top_k/2 - -Top-k values and indices (`%Nx.Block.TopK{}`). - -### Lowering - -Native `Torchx.top_k/2`. - -### Options - -* `:k` — passed to LibTorch - ## fft2/2 Two-dimensional FFT (`%Nx.Block.FFT2{}`). -### Lowering - -Native `Torchx.fft2/3`. - ### Numerical notes Signed zeros and `inspect` formatting may differ from `Nx.BinaryBackend`; see `Torchx.NxBlockTest` for numerical comparisons. -## ifft2/2 - -Two-dimensional inverse FFT (`%Nx.Block.IFFT2{}`). - -### Lowering - -Native `Torchx.ifft2/3`. - -## rfft/2 - -Real FFT (`%Nx.Block.RFFT{}`). - -### Lowering - -Native `Torchx.rfft/3`. - -## irfft/2 - -Inverse real FFT (`%Nx.Block.IRFFT{}`). - -### Lowering - -Native `Torchx.irfft/3`. - -## cumulative_sum/2 - -Cumulative sum (`%Nx.Block.CumulativeSum{}`). - -### Lowering - -Native `Torchx.cumulative_sum/2` with optional `flip` when `:reverse` is true. - -### Options - -* `:axis`, `:reverse` — honoured - -## cumulative_product/2 - -Cumulative product (`%Nx.Block.CumulativeProduct{}`). - -### Lowering - -Native `Torchx.cumulative_product/2`. - -## cumulative_min/2 - -Cumulative minimum (`%Nx.Block.CumulativeMin{}`). - -### Lowering - -Native `Torchx.cumulative_min/2`. - -## cumulative_max/2 - -Cumulative maximum (`%Nx.Block.CumulativeMax{}`). - -### Lowering - -Native `Torchx.cumulative_max/2`. - -## all_close/3 - -All-close comparison (`%Nx.Block.AllClose{}`). - -### Lowering - -Native `Torchx.all_close/4` after dtype merge of operands. - -### Options - -* `:rtol`, `:atol`, `:equal_nan` — passed to LibTorch - -## logical_not/1 - -Logical not (`%Nx.Block.LogicalNot{}`). - -### Lowering - -Native `Torchx.logical_not/1`. - ## phase/1 Complex phase (`%Nx.Block.Phase{}`). -### Lowering +### Behaviour -Default Nx block callback (no dedicated LibTorch wrapper in `Torchx.Backend`). +No dedicated LibTorch wrapper — Torchx invokes the default Nx block callback. ## runtime_call/4 @@ -184,15 +65,6 @@ Transfer tensor data to another backend. `backend_copy/2` followed by `backend_deallocate/1` on the source tensor. -## backend_deallocate/1 - -Deallocate LibTorch tensor storage. - -### Behaviour - -Calls `Torchx.delete_tensor/1`. Returns `:already_deallocated` if the -underlying reference is invalid. - ## to_pointer/2 Export tensor memory as a pointer. @@ -208,14 +80,3 @@ Import tensor memory from a pointer. ### Behaviour **Not supported.** `from_pointer/5` raises on `Torchx.Backend`. - -## to_batched/3 - -Stream tensors in batches. - -### Behaviour - -Splits via `Torchx.split/2` on the leading axis. When `:leftover` is -`:repeat`, concatenates a tail slice so the final chunk matches `batch_size`. -May drop a partial last chunk when the size is not divisible and leftover is -not `:repeat`. diff --git a/torchx/guides/backend_documentation/nx_lin_alg.md b/torchx/guides/backend_documentation/nx_lin_alg.md index ee025b05ee..f5d22a20bc 100644 --- a/torchx/guides/backend_documentation/nx_lin_alg.md +++ b/torchx/guides/backend_documentation/nx_lin_alg.md @@ -7,15 +7,8 @@ iex> Nx.shape(l) {2, 2} ``` -Torchx implementation notes for `Nx.LinAlg`. - -Each function below documents how Torchx handles the corresponding `Nx.LinAlg` -operation when `Torchx.Backend` is the active backend. - -Block-tagged operations are dispatched through `Torchx.Backend.block/4`. When -Torchx does not provide a native LibTorch path for a block, it invokes the -default `Nx.block/4` callback, which composes elementary Nx operations on the -same backend. +Torchx-specific notes for `Nx.LinAlg` where behaviour diverges from the portable +Nx API or from other backends. Several linear algebra blocks are **not** supported natively on Apple MPS (`:mps`). On that device, Torchx falls back to the default Nx block implementation instead @@ -27,19 +20,11 @@ back to MPS. Cholesky decomposition (`%Nx.Block.LinAlg.Cholesky{}`). -### Lowering - -On **CPU** and **CUDA**, Torchx calls `Torchx.cholesky/1` (LibTorch -`torch::linalg_cholesky`). +### Platforms On **MPS**, the default Nx callback is used because Cholesky is not in the native MPS path. -### Types - -Follows Nx promotion to floating point in the public API. LibTorch receives -the tensor's Torch scalar type (`:float`, `:double`, `:complex`, etc.). - ### Numerical notes Results may differ slightly from `Nx.BinaryBackend` due to LibTorch numerics @@ -49,10 +34,7 @@ and rounding (half-to-even vs Elixir's half-away-from-zero elsewhere in Torchx). Linear system solve (`%Nx.Block.LinAlg.Solve{}`). -### Lowering - -On **CPU** and **CUDA**, Torchx calls `Torchx.solve/2` (LibTorch -`torch::linalg_solve`). +### Platforms On **MPS**, the default Nx callback is used. @@ -63,18 +45,11 @@ dependent epsilon (`1.0e-4` for `f32`, `1.0e-10` for `f64`). Singular systems raise `ArgumentError` with `"can't solve for singular matrix"` instead of letting LibTorch throw an opaque error. -### Types - -Operands are cast to a merged floating-point Torch type before the solve. - ## qr/2 QR decomposition (`%Nx.Block.LinAlg.QR{}`). -### Lowering - -Torchx calls `Torchx.qr/2`, mapping `mode: :reduced` to LibTorch's reduced QR -flag. +### Platforms On **MPS**, inputs are copied to CPU, QR is computed there, and `{q, r}` are copied back to MPS. @@ -82,8 +57,7 @@ copied back to MPS. ### Options * `:mode` — `:reduced` or `:complete`, passed to LibTorch -* `:eps` — ignored (used only by the default Nx callback on MPS fallback -paths for other blocks; QR always uses LibTorch when not on MPS) +* `:eps` — ignored (used only by the default Nx callback on MPS fallback paths) ### Numerical notes @@ -94,10 +68,7 @@ via reconstruction `q · r ≈ input`. Hermitian eigendecomposition (`%Nx.Block.LinAlg.Eigh{}`). -### Lowering - -On **CPU** and **CUDA**, integer inputs are promoted to `f32` before calling -`Torchx.eigh/1` (LibTorch `torch::linalg_eigh`). +### Platforms On **MPS**, the default Nx callback is used. @@ -115,15 +86,11 @@ eigenvector equality. Singular value decomposition (`%Nx.Block.LinAlg.SVD{}`). -### Lowering - -Torchx calls `Torchx.svd/2`, passing `full_matrices?` to LibTorch. +### Platforms On **MPS**, inputs are copied to CPU, SVD is computed there, and outputs are copied back to MPS. -Integer inputs are promoted to `f32` before the native call. - ### Types Complex SVD is **not** supported — `Nx.LinAlg.svd/2` raises for complex inputs @@ -143,41 +110,23 @@ Singular values are the most stable quantity to compare across backends. LU decomposition with partial pivoting (`%Nx.Block.LinAlg.LU{}`). -### Lowering - -On **CPU** and **CUDA**, integer inputs are promoted to `f32` before calling -`Torchx.lu/1` (LibTorch `torch::linalg_lu`). +### Platforms On **MPS**, the default Nx callback is used. -### Types - -Permutation `p` keeps integer semantics from LibTorch; `l` and `u` are -floating point. - ## determinant/1 Matrix determinant (`%Nx.Block.LinAlg.Determinant{}`). -### Lowering - -On **CPU** and **CUDA**, integer inputs are promoted to `f32` before calling -`Torchx.determinant/1`. +### Platforms On **MPS**, the default Nx callback is used. -### Types - -Output dtype follows LibTorch / Nx promotion rules for the input. - ## triangular_solve/3 Triangular solve (direct `triangular_solve` callback, not an `Nx.block/4`). -### Lowering - -Torchx calls `Torchx.triangular_solve/4` (LibTorch -`torch::linalg_solve_triangular`). +### Platforms On **MPS**, `a` and `b` are transferred to CPU for the solve; the result is transferred back to MPS. @@ -192,8 +141,3 @@ transferred back to MPS. ### Singular matrices Same determinant-based singularity check as `solve/2` before calling LibTorch. - -### Numerical notes - -Batched inputs are reshaped to `{batch, m, m}` and `{batch, m, n}` internally -when needed. diff --git a/torchx/lib/torchx.ex b/torchx/lib/torchx.ex index e94410e6a3..387560739e 100644 --- a/torchx/lib/torchx.ex +++ b/torchx/lib/torchx.ex @@ -160,9 +160,10 @@ defmodule Torchx do ## Backend documentation - Per-operation implementation notes for Nx linear algebra are documented in - [Nx.LinAlg](backend_documentation-nx_lin_alg.html). Top-level `Nx` blocks and - transfers are documented in [Nx](backend_documentation-nx.html). See the + Backend-specific behaviour, options, and limitations for Nx linear algebra are + documented in [Nx.LinAlg](backend_documentation-nx_lin_alg.html). Top-level + `Nx` blocks and transfers are documented in + [Nx](backend_documentation-nx.html). See the [backend documentation convention](https://hexdocs.pm/nx/backend_documentation-convention.html) in the Nx package. """ diff --git a/torchx/lib/torchx/backend.ex b/torchx/lib/torchx/backend.ex index 53324dc767..8d1e5ede0f 100644 --- a/torchx/lib/torchx/backend.ex +++ b/torchx/lib/torchx/backend.ex @@ -27,7 +27,7 @@ defmodule Torchx.Backend do and nan becomes zero. `Torchx` behaviour is type dependent with no clear rule across types. - Per-operation implementation notes are documented in the + Backend-specific behaviour, options, and limitations are documented in the [Backend documentation](backend_documentation.html) guides — see the [backend documentation convention](https://hexdocs.pm/nx/backend_documentation-convention.html) in the Nx package. diff --git a/torchx/mix.exs b/torchx/mix.exs index 6e842be44f..8efdedc616 100644 --- a/torchx/mix.exs +++ b/torchx/mix.exs @@ -60,16 +60,13 @@ defmodule Torchx.MixProject do extras: [ "CHANGELOG.md", "guides/backend_documentation/index.md": [ - filename: "backend_documentation", - title: "Backend documentation" + filename: "backend_documentation" ], "guides/backend_documentation/nx.md": [ - filename: "backend_documentation-nx", - title: "Nx (Torchx)" + filename: "backend_documentation-nx" ], "guides/backend_documentation/nx_lin_alg.md": [ - filename: "backend_documentation-nx_lin_alg", - title: "Nx.LinAlg (Torchx)" + filename: "backend_documentation-nx_lin_alg" ] ], groups_for_extras: [ From dd75511d39b575174d235ec5e097df74bb5289fd Mon Sep 17 00:00:00 2001 From: Arthur Romeu <14204271+Chapaman@users.noreply.github.com> Date: Tue, 7 Jul 2026 09:02:42 -0300 Subject: [PATCH 4/6] Apply suggestions from code review Co-authored-by: Paulo Valente <16843419+polvalente@users.noreply.github.com> --- exla/guides/backend_documentation/index.md | 9 --------- exla/guides/backend_documentation/nx.md | 8 -------- exla/guides/backend_documentation/nx_lin_alg.md | 7 ------- nx/lib/nx.ex | 4 +--- torchx/guides/backend_documentation/index.md | 10 ---------- torchx/guides/backend_documentation/nx.md | 8 -------- torchx/guides/backend_documentation/nx_lin_alg.md | 7 ------- 7 files changed, 1 insertion(+), 52 deletions(-) diff --git a/exla/guides/backend_documentation/index.md b/exla/guides/backend_documentation/index.md index bd978f1b10..248617fa79 100644 --- a/exla/guides/backend_documentation/index.md +++ b/exla/guides/backend_documentation/index.md @@ -13,12 +13,3 @@ the [backend documentation convention](https://hexdocs.pm/nx/backend_documentati Implementation code lives in `EXLA.Defn`, `EXLA.CustomCall`, and related modules. -## Example - -```elixir -iex> Nx.take(Nx.tensor([10, 20, 30]), Nx.tensor([0, 2])) -#Nx.Tensor< - s32[2] - [10, 30] -> -``` diff --git a/exla/guides/backend_documentation/nx.md b/exla/guides/backend_documentation/nx.md index a1e9d1e00f..a18486de65 100644 --- a/exla/guides/backend_documentation/nx.md +++ b/exla/guides/backend_documentation/nx.md @@ -1,13 +1,5 @@ # Nx (EXLA) -```elixir -iex> Nx.take(Nx.tensor([10, 20, 30]), Nx.tensor([0, 2])) -#Nx.Tensor< - s32[2] - [10, 30] -> -``` - EXLA-specific notes for top-level `Nx` operations where behaviour diverges from the portable Nx API or from other backends. diff --git a/exla/guides/backend_documentation/nx_lin_alg.md b/exla/guides/backend_documentation/nx_lin_alg.md index 7c00d62cdd..a685812bd3 100644 --- a/exla/guides/backend_documentation/nx_lin_alg.md +++ b/exla/guides/backend_documentation/nx_lin_alg.md @@ -1,12 +1,5 @@ # Nx.LinAlg (EXLA) -```elixir -iex> a = Nx.tensor([[4.0, 2.0], [2.0, 3.0]]) -iex> l = Nx.LinAlg.cholesky(a) -iex> Nx.shape(l) -{2, 2} -``` - EXLA-specific notes for `Nx.LinAlg` where behaviour diverges from the portable Nx API or from other backends. diff --git a/nx/lib/nx.ex b/nx/lib/nx.ex index a70fb73d45..ed13e11fc3 100644 --- a/nx/lib/nx.ex +++ b/nx/lib/nx.ex @@ -42,9 +42,7 @@ defmodule Nx do Backend-specific notes for extensible operations (`Nx.block/4`, transfers, and related APIs) are documented in the [backend documentation convention](backend_documentation-convention.html) guide and in each backend's **Backend documentation** pages on HexDocs when a - backend has divergent behaviour or backend-specific options (for example - [EXLA](https://hexdocs.pm/exla/backend_documentation.html) and - [Torchx](https://hexdocs.pm/torchx/backend_documentation.html)). + backend has divergent behaviour or backend-specific option. Continue reading this documentation for an overview of creating, diff --git a/torchx/guides/backend_documentation/index.md b/torchx/guides/backend_documentation/index.md index 80a8a7c36d..a3326b7b11 100644 --- a/torchx/guides/backend_documentation/index.md +++ b/torchx/guides/backend_documentation/index.md @@ -12,13 +12,3 @@ the [backend documentation convention](https://hexdocs.pm/nx/backend_documentati * [Nx.LinAlg](backend_documentation-nx_lin_alg.html) — linear algebra blocks Implementation code lives in `Torchx.Backend` and the `Torchx` NIF bindings. - -## Example - -```elixir -iex> Nx.take(Nx.tensor([10, 20, 30]), Nx.tensor([0, 2])) -#Nx.Tensor< - s32[2] - [10, 30] -> -``` diff --git a/torchx/guides/backend_documentation/nx.md b/torchx/guides/backend_documentation/nx.md index 448136b4b3..eafe8e6eb3 100644 --- a/torchx/guides/backend_documentation/nx.md +++ b/torchx/guides/backend_documentation/nx.md @@ -1,13 +1,5 @@ # Nx (Torchx) -```elixir -iex> Nx.take(Nx.tensor([10, 20, 30]), Nx.tensor([0, 2])) -#Nx.Tensor< - s32[2] - [10, 30] -> -``` - Torchx-specific notes for top-level `Nx` operations where behaviour diverges from the portable Nx API or from other backends. diff --git a/torchx/guides/backend_documentation/nx_lin_alg.md b/torchx/guides/backend_documentation/nx_lin_alg.md index f5d22a20bc..31809b5ab4 100644 --- a/torchx/guides/backend_documentation/nx_lin_alg.md +++ b/torchx/guides/backend_documentation/nx_lin_alg.md @@ -1,12 +1,5 @@ # Nx.LinAlg (Torchx) -```elixir -iex> a = Nx.tensor([[4.0, 2.0], [2.0, 3.0]]) -iex> l = Nx.LinAlg.cholesky(a) -iex> Nx.shape(l) -{2, 2} -``` - Torchx-specific notes for `Nx.LinAlg` where behaviour diverges from the portable Nx API or from other backends. From d2630d673823274b62a5e330482db27332b2543b Mon Sep 17 00:00:00 2001 From: Chapaman <14204271+Chapaman@users.noreply.github.com> Date: Tue, 7 Jul 2026 09:07:35 -0300 Subject: [PATCH 5/6] remove exla&torchx mentions from linalg --- nx/lib/nx/lin_alg.ex | 37 ------------------------------------- 1 file changed, 37 deletions(-) diff --git a/nx/lib/nx/lin_alg.ex b/nx/lib/nx/lin_alg.ex index e115b60dd5..2831cb8a62 100644 --- a/nx/lib/nx/lin_alg.ex +++ b/nx/lib/nx/lin_alg.ex @@ -136,10 +136,6 @@ defmodule Nx.LinAlg do ] ] > - - ## Backend notes - - See [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx_lin_alg.html#cholesky-1). """ def cholesky(tensor) do %T{vectorized_axes: vectorized_axes} = tensor = Nx.to_tensor(tensor) @@ -587,11 +583,6 @@ defmodule Nx.LinAlg do iex> a = Nx.tensor([[1, 0, 0], [1, 1, 0], [1, 1, 1]], type: :f64) iex> Nx.LinAlg.triangular_solve(a, Nx.tensor([1, 2, 1]), transform_a: :other) ** (ArgumentError) invalid value for :transform_a option, expected :none, :transpose, or :conjugate, got: :other - - - ## Backend notes - - See [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx_lin_alg.html#triangular_solve-3). """ def triangular_solve(a, b, opts \\ []) do opts = keyword!(opts, lower: true, left_side: true, transform_a: :none) @@ -705,10 +696,6 @@ defmodule Nx.LinAlg do iex> Nx.LinAlg.solve(Nx.tensor([[3, 0, 0, 0], [2, 1, 0, 0], [1, 1, 1, 1]]), Nx.tensor([4])) ** (ArgumentError) `a` tensor has incompatible dimensions, expected a square matrix or a batch of square matrices, got: {3, 4} - - ## Backend notes - - See [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx_lin_alg.html#solve-2). """ # IMPORTANT: This function cannot be a defn because # optional needs to work on the actual backend. @@ -1135,11 +1122,6 @@ defmodule Nx.LinAlg do iex> t = Nx.tensor([[-3, 2, 1], [0, 1, 1], [0, 0, -1]]) iex> Nx.LinAlg.qr(t, mode: :error_test) ** (ArgumentError) invalid :mode received. Expected one of [:reduced, :complete], received: :error_test - - ## Backend notes - - See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx_lin_alg.html#qr-2) and - [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx_lin_alg.html#qr-2). """ def qr(tensor, opts \\ []) do opts = keyword!(opts, mode: :reduced, eps: 1.0e-10) @@ -1400,11 +1382,6 @@ defmodule Nx.LinAlg do iex> Nx.LinAlg.eigh(Nx.tensor([[1, 2, 3], [4, 5, 6]])) ** (ArgumentError) tensor must be a square matrix or a batch of square matrices, got shape: {2, 3} - - ## Backend notes - - See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx_lin_alg.html#eigh-2) and - [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx_lin_alg.html#eigh-2). """ def eigh(tensor, opts \\ []) do opts = keyword!(opts, max_iter: 1_000, eps: 1.0e-4) @@ -1528,11 +1505,6 @@ defmodule Nx.LinAlg do [0.0, 0.0, 1.0] ] > - - ## Backend notes - - See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx_lin_alg.html#svd-2) and - [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx_lin_alg.html#svd-2). """ def svd(tensor, opts \\ []) do opts = keyword!(opts, max_iter: 100, full_matrices?: true) @@ -1763,10 +1735,6 @@ defmodule Nx.LinAlg do iex> Nx.LinAlg.lu(Nx.tensor([[1, 1, 1, 1], [-1, 4, 4, -1], [4, -2, 2, 0]])) ** (ArgumentError) tensor must be a square matrix or a batch of square matrices, got shape: {3, 4} - - ## Backend notes - - See [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx_lin_alg.html#lu-1). """ def lu(tensor) do %T{vectorized_axes: vectorized_axes} = tensor = Nx.to_tensor(tensor) @@ -2015,11 +1983,6 @@ defmodule Nx.LinAlg do c64 -0.0-6.0i > - - - ## Backend notes - - See [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx_lin_alg.html#determinant-1). """ # IMPORTANT: This function cannot be a defn because # optional needs to work on the actual backend. From 7a37a57085a771f128d25b04730d80979f2e8cab Mon Sep 17 00:00:00 2001 From: Chapaman <14204271+Chapaman@users.noreply.github.com> Date: Fri, 10 Jul 2026 20:40:42 -0300 Subject: [PATCH 6/6] update according to feedback --- .../backend_documentation/nx_lin_alg.md | 15 ++++---- exla/mix.exs | 3 +- nx/lib/nx.ex | 35 ------------------- .../backend_documentation/nx_lin_alg.md | 20 +++++++---- 4 files changed, 23 insertions(+), 50 deletions(-) diff --git a/exla/guides/backend_documentation/nx_lin_alg.md b/exla/guides/backend_documentation/nx_lin_alg.md index a685812bd3..d3e1ec8bce 100644 --- a/exla/guides/backend_documentation/nx_lin_alg.md +++ b/exla/guides/backend_documentation/nx_lin_alg.md @@ -21,15 +21,18 @@ complex inputs and unsupported types ### Options -* `:mode` — honoured by whichever implementation is selected (`:reduced` or +* `:mode` — honored by whichever implementation is selected (`:reduced` or `:complete`) * `:eps` — used by the default callback only; ignored by the native custom call ### Numerical notes Native and default implementations may differ slightly in orthogonality of `q` -and triangular structure of `r` within floating-point tolerance. Reconstruction -`q · r` should match the input within reasonable tolerances for the dtype. +and triangular structure of `r` within floating-point tolerance. +For example, for an eigenpair `(lambda, v)`, `(-lambda, -v)` is +an equally valid eigenpair that would lead to radically different +`q` and `r` matrices. Reconstruction of `q · r` should match the +input within reasonable tolerances for the dtype. ## eigh/2 @@ -44,7 +47,7 @@ unsupported types ### Options -* `:max_iter` and `:eps` — honoured by the default callback; ignored by the +* `:max_iter` and `:eps` — honored by the default callback; ignored by the native custom call ### Numerical notes @@ -59,8 +62,8 @@ Singular value decomposition (`%Nx.Block.LinAlg.SVD{}`). ### Options -* `:max_iter` — honoured (default iterative algorithm in Nx) -* `:full_matrices?` — honoured +* `:max_iter` — honored (default iterative algorithm in Nx) +* `:full_matrices?` — honored ### Numerical notes diff --git a/exla/mix.exs b/exla/mix.exs index 59c06b0473..2db1588fe8 100644 --- a/exla/mix.exs +++ b/exla/mix.exs @@ -70,8 +70,7 @@ defmodule EXLA.MixProject do defp deps do [ - {:nx, "~> 0.12.0"}, - # {:nx, path: "../nx"}, + {:nx, path: "../nx"}, {:telemetry, "~> 0.4.0 or ~> 1.0"}, {:xla, "~> 0.10.0", runtime: false}, {:fine, "~> 0.1", runtime: false}, diff --git a/nx/lib/nx.ex b/nx/lib/nx.ex index ed13e11fc3..9af0bc2722 100644 --- a/nx/lib/nx.ex +++ b/nx/lib/nx.ex @@ -2277,11 +2277,6 @@ defmodule Nx do Inside `defn`, this builds an expression node that is supported by compilers. Outside of `defn` or in backends without special support, it executes `fun` directly and validates that the result matches the provided template. - - ## Backend notes - - See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#runtime_call-4) and - [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#runtime_call-4). """ @doc type: :backend def runtime_call(output, tensor_or_container, opts \\ [], fun) when is_function(fun, 2) do @@ -4991,12 +4986,6 @@ defmodule Nx do [4, 5, 6] ] > - - - ## Backend notes - - See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#backend_copy-2) and - [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#backend_copy-2). """ @doc type: :backend def backend_copy(tensor_or_container, backend \\ Nx.BinaryBackend) do @@ -5046,12 +5035,6 @@ defmodule Nx do Transfer the device tensor back to an Elixir tensor: tensor = Nx.backend_transfer(device_tensor) - - - ## Backend notes - - See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#backend_transfer-2) and - [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#backend_transfer-2). """ @doc type: :backend def backend_transfer(tensor_or_container, backend \\ Nx.BinaryBackend) do @@ -8714,10 +8697,6 @@ defmodule Nx do f32[2] [1.1071488, 2.6779451] > - - ## Backend notes - - See [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#phase-1). """ @doc type: :element def phase(tensor) do @@ -16936,10 +16915,6 @@ defmodule Nx do iex> Nx.fft2(Nx.tensor([1, 1]), length: :invalid) ** (ArgumentError) expected a tensor with rank > 1, got tensor with rank 1 - - ## Backend notes - - See [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#fft2-2). """ @doc type: :ndim def fft2(tensor, opts \\ []), do: call_fft2(tensor, opts, :fft2) @@ -17670,11 +17645,6 @@ defmodule Nx do #=> [10, 20, 30] #=> ] #=> > - - ## Backend notes - - See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#from_pointer-5) and - [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#from_pointer-5). """ @doc type: :creation def from_pointer(backend, pointer, type, shape, opts \\ []) @@ -17720,11 +17690,6 @@ defmodule Nx do t = Nx.s32([1, 2, 3]) Nx.to_pointer(t, mode: :ipc) #=> %Nx.Pointer{kind: :ipc, address: nil, data_size: 32, handle: "some-ipc-handle"} - - ## Backend notes - - See [EXLA](https://hexdocs.pm/exla/backend_documentation-nx.html#to_pointer-2) and - [Torchx](https://hexdocs.pm/torchx/backend_documentation-nx.html#to_pointer-2). """ @doc type: :creation def to_pointer(tensor, opts \\ []) do diff --git a/torchx/guides/backend_documentation/nx_lin_alg.md b/torchx/guides/backend_documentation/nx_lin_alg.md index 31809b5ab4..2032f83916 100644 --- a/torchx/guides/backend_documentation/nx_lin_alg.md +++ b/torchx/guides/backend_documentation/nx_lin_alg.md @@ -54,8 +54,12 @@ copied back to MPS. ### Numerical notes -`q` and `r` are not guaranteed to match other backends bit-for-bit. Verify -via reconstruction `q · r ≈ input`. +Results may differ slightly in orthogonality of `q` and triangular structure +of `r` within floating-point tolerance. +For example, for an eigenpair `(lambda, v)`, `(-lambda, -v)` is +an equally valid eigenpair that would lead to radically different +`q` and `r` matrices. Reconstruction of `q · r` should match the +input within reasonable tolerances for the dtype. ## eigh/2 @@ -67,13 +71,14 @@ On **MPS**, the default Nx callback is used. ### Options -* `:max_iter` and `:eps` — honoured only by the default Nx callback (MPS +* `:max_iter` and `:eps` — honored only by the default Nx callback (MPS path); ignored by the LibTorch kernel ### Numerical notes -Eigenvectors are not unique. Prefer comparing reconstructions over direct -eigenvector equality. +Eigenvectors are not unique (sign and degenerate subspaces). Compare results +using reconstruction `v · diag(λ) · vᵀ` rather than direct eigenvector equality +across backends. ## svd/2 @@ -91,12 +96,13 @@ on Torchx. ### Options -* `:full_matrices?` — honoured by LibTorch +* `:full_matrices?` — honored by LibTorch * `:max_iter` — ignored (LibTorch uses its own algorithm; the default Nx iterative implementation is not used on the native path) ### Numerical notes +SVD is iterative; results may vary across devices and precision settings. Singular values are the most stable quantity to compare across backends. ## lu/1 @@ -126,7 +132,7 @@ transferred back to MPS. ### Options -* `:lower` — honoured (`upper` flag inverted for LibTorch) +* `:lower` — honored (`upper` flag inverted for LibTorch) * `:transform_a` — `:none` and `:transpose` are supported * `:left_side` — **only `true` is supported**; `left_side: false` raises `ArgumentError`