Standardize docs across nx, exla, torchx

exla/guides/backend_documentation/index.md

@@ -0,0 +1,15 @@
+# Backend documentation
+
+EXLA-specific documentation for Nx backend behaviour.
+
+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
+
+Implementation code lives in `EXLA.Defn`, `EXLA.CustomCall`, and related modules.
+

exla/guides/backend_documentation/nx.md

@@ -0,0 +1,92 @@
+# Nx (EXLA)
+
+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).
+
+## runtime_call/4
+
+Runtime Elixir callback from `defn` (`runtime_call` expression).
+
+### Platforms
+
+* **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`
+
+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`.
+
+### Options
+
+* `:client`, `:device_id` — select the EXLA client and device
+
+## 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
+
+### Limitations
+
+Not supported for ROCm or TPU.

exla/guides/backend_documentation/nx_lin_alg.md

@@ -0,0 +1,72 @@
+# Nx.LinAlg (EXLA)
+
+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`.
+
+## qr/2
+
+QR decomposition (`%Nx.Block.LinAlg.QR{}`).
+
+### Supported platforms
+
+* **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` — 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.
+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
+
+Hermitian eigendecomposition (`%Nx.Block.LinAlg.Eigh{}`).
+
+### Supported platforms
+
+* **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
+
+* `:max_iter` and `:eps` — honored 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{}`).
+
+### Options
+
+* `:max_iter` — honored (default iterative algorithm in Nx)
+* `:full_matrices?` — honored
+
+### 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.

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.
 
+  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.
+
   ## Clients
 
   The `EXLA` library uses a client for compiling and executing code.

exla/mix.exs

@@ -87,9 +87,21 @@ 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"
+        ],
+        "guides/backend_documentation/nx.md": [
+          filename: "backend_documentation-nx"
+        ],
+        "guides/backend_documentation/nx_lin_alg.md": [
+          filename: "backend_documentation-nx_lin_alg"
+        ]
       ],
       skip_undefined_reference_warnings_on: ["CHANGELOG.md"],
+      groups_for_extras: [
+        "Backend documentation": ~r"^guides/backend_documentation/"
+      ],
       groups_for_modules: [
         # EXLA,
         # EXLA.Backend,

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

nx/guides/backend_documentation/convention.md

@@ -0,0 +1,55 @@
+# 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 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. 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.
+
+## What to document
+
+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
+
+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 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 divergent behaviour for 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]
+>
+```

nx/lib/nx.ex

@@ -39,6 +39,12 @@ 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 when a
+  backend has divergent behaviour or backend-specific option.
+
+
   Continue reading this documentation for an overview of creating,
   broadcasting, and accessing/slicing Nx tensors.
 
@@ -2515,6 +2521,7 @@ defmodule Nx do
         s32
         [3, 1]
       >
+
   """
   @doc type: :conversion
   def to_batched(tensor, batch_size, opts \\ [])
@@ -4979,7 +4986,6 @@ defmodule Nx do
           [4, 5, 6]
         ]
       >
-
   """
   @doc type: :backend
   def backend_copy(tensor_or_container, backend \\ Nx.BinaryBackend) do
@@ -5029,7 +5035,6 @@ defmodule Nx do
   Transfer the device tensor back to an Elixir tensor:
 
       tensor = Nx.backend_transfer(device_tensor)
-
   """
   @doc type: :backend
   def backend_transfer(tensor_or_container, backend \\ Nx.BinaryBackend) do
@@ -5049,6 +5054,7 @@ defmodule Nx do
   It returns either `:ok` or `:already_deallocated`.
 
   Note: This function cannot be used in `defn`.
+
   """
   @doc type: :backend
   def backend_deallocate(tensor_or_container) do
@@ -7223,6 +7229,7 @@ defmodule Nx do
         [0, 1, 0]
       >
 
+
   """
   @doc type: :element
   def logical_not(tensor) do
@@ -9276,6 +9283,7 @@ defmodule Nx do
           [0, 0]
         ]
       >
+
   """
   @doc type: :aggregation
   def all_close(a, b, opts \\ []) do
@@ -11680,6 +11688,7 @@ defmodule Nx do
           [2, 3, 6]
         ]
       >
+
   """
   @doc type: :cumulative
   def cumulative_sum(tensor, opts \\ []),
@@ -11756,6 +11765,7 @@ defmodule Nx do
           [2, 2, 6]
         ]
       >
+
   """
   @doc type: :cumulative
   def cumulative_product(tensor, opts \\ []),
@@ -11832,6 +11842,7 @@ defmodule Nx do
           [2, 1, 1]
         ]
       >
+
   """
   @doc type: :cumulative
   def cumulative_min(tensor, opts \\ []),
@@ -11908,6 +11919,7 @@ defmodule Nx do
           [2, 2, 3]
         ]
       >
+
   """
   @doc type: :cumulative
   def cumulative_max(tensor, opts \\ []),
@@ -15667,6 +15679,7 @@ defmodule Nx do
       iex> Nx.top_k(a, k: 1)
       ** (ArgumentError) top_k input must have at least rank 1
 
+
   """
   @doc type: :ndim
   def top_k(tensor, opts \\ []) do

nx/lib/nx/backend.ex

@@ -28,6 +28,11 @@ 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 divergent behaviour, backend-specific options, and limitations — see
+  the [backend documentation convention](backend_documentation-convention.html)
+  guide.
   """
 
   @type t :: %{__struct__: atom()}