jax.jit missing manual ?

[mitkotak]

I am wondering what the equivalent of this is in JAX-land ?

I am a big fan of JAX docs (particularly autodidax) but I am wondering if there's a similar way to reason through the entire stack all the way down to the Triton kernels that XLA is generating ?

[yashk2810]

JAX has jax.jit for automatic partitioning: https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html

and shard_map for manual partitioning: https://jax.readthedocs.io/en/latest/notebooks/shard_map.html

If you want to write kernels, check out pallas: https://jax.readthedocs.io/en/latest/pallas/quickstart.html