diff --git a/exla/guides/backend_documentation/index.md b/exla/guides/backend_documentation/index.md new file mode 100644 index 0000000000..248617fa79 --- /dev/null +++ b/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. + diff --git a/exla/guides/backend_documentation/nx.md b/exla/guides/backend_documentation/nx.md new file mode 100644 index 0000000000..a18486de65 --- /dev/null +++ b/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. 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..d3e1ec8bce --- /dev/null +++ b/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. diff --git a/exla/lib/exla.ex b/exla/lib/exla.ex index 9be5dc13db..5b019a57d9 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. + 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. diff --git a/exla/mix.exs b/exla/mix.exs index 035450b640..2db1588fe8 100644 --- a/exla/mix.exs +++ b/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, diff --git a/exla/test/exla/backend_documentation_test.exs b/exla/test/exla/backend_documentation_test.exs new file mode 100644 index 0000000000..081319db10 --- /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..d21cabe27a --- /dev/null +++ b/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] +> +``` diff --git a/nx/lib/nx.ex b/nx/lib/nx.ex index 77dd527889..9af0bc2722 100644 --- a/nx/lib/nx.ex +++ b/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 diff --git a/nx/lib/nx/backend.ex b/nx/lib/nx/backend.ex index 646356228a..7d8c938786 100644 --- a/nx/lib/nx/backend.ex +++ b/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()} diff --git a/nx/lib/nx/lin_alg.ex b/nx/lib/nx/lin_alg.ex index e7e166c6ab..2831cb8a62 100644 --- a/nx/lib/nx/lin_alg.ex +++ b/nx/lib/nx/lin_alg.ex @@ -583,7 +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 - """ def triangular_solve(a, b, opts \\ []) do opts = keyword!(opts, lower: true, left_side: true, transform_a: :none) @@ -1984,7 +1983,6 @@ defmodule Nx.LinAlg do c64 -0.0-6.0i > - """ # 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..f6c3c24e22 100644 --- a/nx/mix.exs +++ b/nx/mix.exs @@ -74,7 +74,10 @@ 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" + ] ], skip_undefined_reference_warnings_on: ["CHANGELOG.md"], groups_for_docs: [ @@ -132,6 +135,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..f90119a44d --- /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..a3326b7b11 --- /dev/null +++ b/torchx/guides/backend_documentation/index.md @@ -0,0 +1,14 @@ +# Backend documentation + +Torchx-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 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 new file mode 100644 index 0000000000..eafe8e6eb3 --- /dev/null +++ b/torchx/guides/backend_documentation/nx.md @@ -0,0 +1,74 @@ +# Nx (Torchx) + +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). + +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. + +## fft2/2 + +Two-dimensional FFT (`%Nx.Block.FFT2{}`). + +### Numerical notes + +Signed zeros and `inspect` formatting may differ from `Nx.BinaryBackend`; see +`Torchx.NxBlockTest` for numerical comparisons. + +## phase/1 + +Complex phase (`%Nx.Block.Phase{}`). + +### Behaviour + +No dedicated LibTorch wrapper — Torchx invokes the default Nx block callback. + +## 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. + +## 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`. 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..2032f83916 --- /dev/null +++ b/torchx/guides/backend_documentation/nx_lin_alg.md @@ -0,0 +1,142 @@ +# Nx.LinAlg (Torchx) + +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 +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{}`). + +### Platforms + +On **MPS**, the default Nx callback is used because Cholesky is not in the +native MPS path. + +### 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{}`). + +### Platforms + +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. + +## qr/2 + +QR decomposition (`%Nx.Block.LinAlg.QR{}`). + +### Platforms + +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) + +### Numerical notes + +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 + +Hermitian eigendecomposition (`%Nx.Block.LinAlg.Eigh{}`). + +### Platforms + +On **MPS**, the default Nx callback is used. + +### Options + +* `:max_iter` and `:eps` — honored only by the default Nx callback (MPS +path); ignored by the LibTorch kernel + +### 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{}`). + +### Platforms + +On **MPS**, inputs are copied to CPU, SVD is computed there, and outputs are +copied back to MPS. + +### Types + +Complex SVD is **not** supported — `Nx.LinAlg.svd/2` raises for complex inputs +on Torchx. + +### Options + +* `: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 + +LU decomposition with partial pivoting (`%Nx.Block.LinAlg.LU{}`). + +### Platforms + +On **MPS**, the default Nx callback is used. + +## determinant/1 + +Matrix determinant (`%Nx.Block.LinAlg.Determinant{}`). + +### Platforms + +On **MPS**, the default Nx callback is used. + +## triangular_solve/3 + +Triangular solve (direct `triangular_solve` callback, not an `Nx.block/4`). + +### Platforms + +On **MPS**, `a` and `b` are transferred to CPU for the solve; the result is +transferred back to MPS. + +### Options + +* `: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` + +### Singular matrices + +Same determinant-based singularity check as `solve/2` before calling LibTorch. diff --git a/torchx/lib/torchx.ex b/torchx/lib/torchx.ex index 506fa958ec..387560739e 100644 --- a/torchx/lib/torchx.ex +++ b/torchx/lib/torchx.ex @@ -157,6 +157,15 @@ defmodule Torchx do PyTorch implements a variety of devices, which can be seen below. #{@valid_devices_md_list} + + ## Backend documentation + + 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. """ use Torchx.Macro alias Torchx.NIF diff --git a/torchx/lib/torchx/backend.ex b/torchx/lib/torchx/backend.ex index 3bbc195b0e..8d1e5ede0f 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. + 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. + ## Options * `:device` - Defaults to `Torchx.default_device/0`. An atom representing the diff --git a/torchx/mix.exs b/torchx/mix.exs index d53c27bc34..8efdedc616 100644 --- a/torchx/mix.exs +++ b/torchx/mix.exs @@ -58,7 +58,19 @@ 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" + ], + "guides/backend_documentation/nx.md": [ + filename: "backend_documentation-nx" + ], + "guides/backend_documentation/nx_lin_alg.md": [ + filename: "backend_documentation-nx_lin_alg" + ] + ], + groups_for_extras: [ + "Backend documentation": ~r"^guides/backend_documentation/" ] ] end @@ -68,7 +80,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", "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..ed0140fac5 --- /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