JAX profiling tools support

google / jax

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more

http://jax.readthedocs.io/

Apache License 2.0

29.06k stars 2.66k forks source link

Open PrajP opened 1 month ago

PrajP commented 1 month ago

Hi Team, Can you please add - Profiling tools that expose performance information at some within-kernel granularity?

Thank you, Prajakta.

sharadmv commented 1 month ago

Are you referring to Pallas/Mosaic/Triton kernels? Or to XLA generated kernels as well?