5 Commits

Author SHA1 Message Date
Skye Wanderman-Milne
b68a08adf1 Add programmatic profiling APIs, and rename some existing APIs.
This change provides aliases for the renamed APIs so existing code
won't break. We should remove these aliases after the next release.
2021-03-25 11:50:55 -07:00
Peter Hawkins
fb7b7dc605 [JAX] Move implementation of profiler.py into jax._src.
PiperOrigin-RevId: 340660150
2020-11-04 09:06:40 -08:00
Jiho Choi
e073e25427
Add a jax.profiler.StepTraceContext API. (#3591)
* Add RootTraceContext.

* Rename to StepTraceContext.
2020-07-10 17:20:06 -07:00
Peter Hawkins
5116fd47aa
Add a heap profiler API and document it. (#3576) 2020-06-26 17:09:09 -04:00
Peter Hawkins
ddd76b803b
Add a jax.profiler module that exposes the new TensorFlow profiler in… (#2354)
* Add a jax.profiler module that exposes the new TensorFlow profiler integration.

* Fix documentation.
2020-03-05 12:07:57 -05:00