iree-org
[New documentation] Write docs for "JAX frontend"
- Live preview of this page: https://scotttodd.github.io/iree/frontends/jax/, which will eventually live on https://google.github.io/iree/ after we switch over the docs.
- ~~Source currently lives on the
staging-docs-reworkbranch here: https://github.com/google/iree/blob/staging-docs-rework/mkdocs/docs/frontends/jax.md . PRs or direct commits to the branch are okay. We'll want to be careful about force pushes / rebases though.~~ Source lives in https://github.com/google/iree/blob/main/docs/website/docs/ml-frameworks/jax.md - See the README here for instructions on previewing changes locally: ~~https://github.com/google/iree/tree/staging-docs-rework/mkdocs~~ https://github.com/google/iree/tree/main/docs/website
- mkdocs + mkdocs material offer a bunch of nice features, like https://squidfunk.github.io/mkdocs-material/reference/code-blocks/ (line numbers, highlighting specific lines, etc.)
Site and doc layout can be changed, but I've sketched out a rough draft already. This file just currently has some ideas in it, but it should match up with the TensorFlow and other frontend docs (see https://github.com/google/iree/issues/5453):
https://github.com/google/iree/blob/97d4eeac924177a094110984a825d8cce6604fa8/mkdocs/docs/frontends/jax.md#L1-L11
I was thinking we could have each "frontend" page discuss the formats used, features supported, and APIs available, while covering general Python binding details in the "bindings/python" page (preview: https://scotttodd.github.io/iree/bindings/python/). Some overlap and linking back and forth is fine, but we don't want people to get too lost in the docs.
Really interested in this - does there exist an update on some branch to play around with, and possibly fill in a few gaps in the docs? :)
The docs in IREE itself are still missing (https://google.github.io/iree/getting-started/jax/, https://github.com/google/iree/blob/main/docs/website/docs/getting-started/jax.md).
The https://github.com/google/iree-jax repo has the latest information and code.
@rsuderman / @stellaraccident , how stable is iree-jax? Could we fill in those docs now?
Starting to work on this, sketches at https://github.com/ScottTodd/iree/tree/docs-jax
https://openxla.org/stablehlo/tutorials/jax-export would be good to link to