11 Commits

Author SHA1 Message Date
Jake VanderPlas
26f2f97805 Document why 'import name as name' is used 2022-12-14 15:07:04 -08:00
Peter Hawkins
ba557d5e1b Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
See https://opensource.google/documentation/reference/releasing/contributions#copyright for more details.

PiperOrigin-RevId: 476167538
2022-09-22 12:27:19 -07:00
Jake VanderPlas
2543542fa8 jax.profiler: remove deprecated functions 2022-07-19 08:13:44 -07:00
Sharad Vikram
449da304b3 Store profiler server as a global variable and add a stop_server function 2022-06-01 17:50:06 -07:00
Jake VanderPlas
5782210174 CI: fix flake8 ignore declarations 2022-04-21 13:44:12 -07:00
Jake VanderPlas
245581411e Add PEP484-compatible export for jax and its subpackages 2021-09-13 14:08:48 -07:00
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