[New documentation] Write docs for "JAX frontend"

ScottTodd commented 3 years ago

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-rework branch 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.

ludgerpaehler commented 2 years ago

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? :)

ScottTodd commented 2 years ago

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?

ScottTodd commented 1 year ago

Starting to work on this, sketches at https://github.com/ScottTodd/iree/tree/docs-jax

ScottTodd commented 2 days ago

https://openxla.org/stablehlo/tutorials/jax-export would be good to link to

iree-org / iree

[New documentation] Write docs for "JAX frontend" #5454