Auto-generated API docs with sphinx

[SamuelMarks]

This PR adds generated API docs.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• [x] I have performed a self-review of my code.
• [x] I have necessary comments in my code, particularly in hard-to-understand areas.
• [x] I have run end-to-end tests tests and provided workload links above if applicable.
• [x] I have made or will make corresponding changes to the doc if needed.

[SamuelMarks]

Looks to work now, I've merged in the #2144 first but can reduce this to one commit—with 10 files changed and you just ignore that it won't work until that is merged first:

Finally, readthedocs deployment is needed. Not sure how you want this done, e.g., GitHub Actions [with a possible security issue if we can't be granular here; preventing overriding of deploy only on main] or out-of-band. Let me know.

Thanks and looking forward to your review

[SamuelMarks]

For my understanding, where do these generated API docs end up? Do we run this manually? Or will it automatically be uploaded to ReadTheDocs?

We can run it on CI, or preferably out-of-band, then that will be loaded into ReadTheDocs.

This is all sphinx now no mkdocs; so is fully integrated with existing sphinx commands like sphinx-build -M html docs out as you can see in my "DOCS.md" (also in this PR).

[SamuelMarks]

Hi - I'm at Quansight Labs and new to the project, and was asked to take a look

I tried building this locally, and can confirm that

First install the dependencies:

$ python3 -m pip install -r requirements_docs.txt

Build

$ sphinx-build -M html docs out

Serve

You can use any static file HTTP server, e.g.:

$ python3 -m http.server -d 'out/html'

works fine for me. The generated API reference in the sidebar looks fine

However, if I click on any API reference item, all I see is something like

Is this intended, or is the toplevel docstring in that module

https://github.com/AI-Hypercomputer/maxtext/blob/28e5097ac467ed8b1d17676d68aa5acc50f9d60d/src/MaxText/accelerator_to_spec_map.py#L15-L19

meant to show up?

Related to that, on the benchmark_chunked_prefill page, I don't see members listed, are we expecting

https://github.com/AI-Hypercomputer/maxtext/blob/28e5097ac467ed8b1d17676d68aa5acc50f9d60d/src/MaxText/benchmark_chunked_prefill.py#L101-L116

to show up?

@MarcoGorelli Thanks for following up. I plan on taking the codebase and going module by module "Gemini: add missing file, module, class, and function docstrings in Google docstring format". Alternatively others can do the same. It should be trivial to maintain 100% doc coverage for this codebase.

As for the specific issues you are coming across, I'll try and replicate. Sometimes sphinx rst's format needs to be explicitly overriden in favour of Google's docstring format; which may be the cause of the issue.

For my understanding, where do these generated API docs end up? Do we run this manually? Or will it automatically be uploaded to ReadTheDocs?

@bvandermoon Thanks for following up. The documentation is auto generated onto readthedocs.

[MarcoGorelli]

I plan on taking the codebase and going module by module

Sure, if it was expected that with this PR then the listed members wouldn't show up, then no objections

Other than that, there's some commented-out code that's being added here, is that intended?