Profiling API's

[wsmoses]

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