Standardize docs across nx, exla, torchx

[Chapaman]

Closes #1759

• Add Nx.Docs to define the convention for backend-specific documentation modules (Kernel.SpecialForms-style ghost modules that mirror the Nx API).
• Add EXLA.Docs.Nx, EXLA.Docs.Nx.LinAlg, Torchx.Docs.Nx, and Torchx.Docs.Nx.LinAlg documenting how each backend lowers or executes block-backed and backend-specific operations.
• Document all current Nx.block/4 entry points: top-level Nx (take, FFT, cumulative ops, transfers, runtime_call, etc.) and Nx.LinAlg (cholesky, QR, SVD, …).
• Add ## Backend notes sections on 29 Nx / Nx.LinAlg functions linking to the EXLA and Torchx docs modules.
• Wire discoverability through Nx, Nx.LinAlg, Nx.Backend, Nx.block/4, EXLA, and Torchx moduledocs, plus Hex “Backend documentation” groups in each package.

[josevalim]

Apologies but I am not a fan of this approach. We are introducing a bunch of modules, which have a runtime cost, to document APIs that do not really exist and people should not invoke. Otherwise this will appear as suggestions on LSP, ExDoc autocompletes, etc.

We should figure out how to structure this using regular pages/guides as much as possible.

[polvalente]

Apologies but I am not a fan of this approach. We are introducing a bunch of modules, which have a runtime cost, to document APIs that do not really exist and people should not invoke. Otherwise this will appear as suggestions on LSP, ExDoc autocompletes, etc.

We should figure out how to structure this using regular pages/guides as much as possible.

I think we can keep the same general idea where we have a standardized nested structure that corresponds to the Nx modules (e.g. EXLA > Backend Documentation > Nx > LinAlg) which is written as Markdown, Livemd or cheatsheets, as relevant. The most important part of this PR for me is having an easy way to look things up structurally.

I suspect livebooks will provide the closest experience to regular modules, and cheatsheets can be complimentary documentation to that. We might also want some sort of CI check to ensure the documentation gets updated if we add or remove functions to the original Nx modules, but we can leave this as a follow-up.

[josevalim]

I think you can pass files to doctest to make sure their content is correct. But that means using .cheatmd or .md files. Any should be fine!

[josevalim]

This is wrong now, as it is only markdown files.

[josevalim]

Question: what documenting this particular behaviour is giving us? What it becomes in StableHLO is not necessarily relevant. If it is, we will probably benefit more from having a debug_stablehlo option that prints it.

[polvalente]

We do have this in EXLA.to_mlir_module

[polvalente]

Now to think out loud about lowering patterns, this function is implemented as an Nx.Block, so. users might want to override the imolementation.

Though at that point they might as well look at the source code.

Not to decide anything, but maybe. what we want here is to scope documenting divergent behavior and backend specific options only?

[josevalim]

Divergent behavior or specific trade-offs. Yes. Most of what EXLA does is already abstracted several layers away, saying we use gather won’t add much. But saying function X does not support A option or has a special flag can help. The question is what users of this library may want to know. If you are implementing custom blocks, then the source is already the best source of truth anyway.

[josevalim]

This is already the title, do we need to give it again? Same above.

[polvalente]

Overall this looks great, I just think we need to clean up some of the things I pointed out

[polvalente]

I don't think we should reference backends from Nx itself