Enzyme-JAX
Enzyme-JAX copied to clipboard
Profiling API's
This https://jax.readthedocs.io/en/latest/profiling.html
goes to https://github.com/jax-ml/jax/blob/0093ba29d898e8f9829e5365485c1f0094f6b0c5/jax/_src/profiler.py#L128
which is implemented here
https://github.com/openxla/xla/blob/c77c7c2c5922b83f219a650a24570f91ffe19f15/xla/python/profiler.cc#L170
Getting api's with this connected should probably fix issues for both reactant and the egraph side, if you guys have a chance to bind/test out these api's
cc @avik-pal @mofeing @aryavohra @smjleo