mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity
16,831 Commits 506 Branches 263 Tags
Commit Graph

26 Commits

Author SHA1 Message Date
Skye Wanderman-Milne
11ff93951c Bump ProfilerTest.test_remote_profiler to run computations for 10 seconds. 
I hate to make a unit test run for this long, but I'm still seeing
cases where the remote profiler will miss the computations and hang
forever. (Why is it so slow?) This should hopefully be enough that it always catches some
computations.
 2023-07-27 23:54:42 +00:00
Skye Wanderman-Milne
2ca151ef5b profiler_test.py fixes and add coverage to Cloud TPU CI 
* Add deps to test requirements, including in new
  `collect-profile-requirements.txt` (to avoid adding tensorflow to
  `test-requirements.txt`).
* Use correct Python executable `ProfilerTest.test_remote_profiler`
  (`python` sometimes defaults to python2)
* Run computations for longer in `ProfilerTest.test_remote_profiler`,
  othewise `collect_profile` sometimes misses it.
 2023-06-20 22:25:17 +00:00
Jake VanderPlas
fbe4f10403 Change to simpler import for jax.config 2023-04-21 11:51:22 -07: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
jax authors
840c96692e Internal change 
PiperOrigin-RevId: 468509799
 2022-08-18 11:39:07 -07:00
Peter Hawkins
1e241dcf16 Catch ModuleNotFoundError instead of ImportError. 
We frequently use the pattern
try:
  import m
except ImportError:
  # do something else.

This suppresses errors when the module can be found but does not import
successfully for any reason. Instead, catch only ModuleNotFoundError so
missing modules are allowed but buggy modules still report errors.
 2022-08-18 15:22:49 +00:00
Sharad Vikram
143ed40a78 Add collect_profile script 2022-06-03 17:56:17 -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
Jeppe Klitgaard
17de89b16a feat: refactor code using pyupgrade 
This PR upgrades legacy Python code to 3.7+ code using pyupgrade:
```sh
pyupgrade --py37-plus --keep-runtime-typing **.py
```

a
 2022-05-17 22:14:05 +01:00
Peter Hawkins
a288b154b4 Fix keyword argument confusion in jax.profiler.annotate_function decorator. 2021-12-06 09:10:44 -05:00
Peter Hawkins
db2e91eba2 Move jax.test_util to jax._src.test_util. 
Add forwarding shims for names used by external clients of JAX in practice.

PiperOrigin-RevId: 398721725
 2021-09-24 07:02:49 -07:00
Peter Hawkins
3ddcec27f2 Update minimum jaxlib version to 0.1.69. 2021-07-15 17:00:13 -04:00
jax authors
3c362ced90 Merge pull request #6380 from skye:python_trace_test 
PiperOrigin-RevId: 368484288
 2021-04-14 12:29:16 -07:00
Peter Hawkins
c8184a87a5 Ensure that profiler tests shut down the profiler, and disable a profiler test. 
May help with CI OOM problems.
 2021-04-12 19:12:30 -04:00
Skye Wanderman-Milne
6d93fa5526 Test that Python tracing is enabled for programmatic profiling 2021-04-08 17:01:13 +00: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
Matthew Johnson
0b76854be0 fix skip logic for test that requires portpicker 2020-11-08 14:48:14 -08:00
Qiao Zhang
42591e9b50 Skip profiler_test if tf_profiler not imported. 2020-11-03 12:50:39 -08:00
Qiao Zhang
1ba59c7310 Fix errors in tests/profiler_test.py. 2020-11-02 12:09:20 -08:00
Qiao Zhang
377f09c570 Add test for profiler client/server interaction. 2020-10-29 17:27:38 -07:00
Jake VanderPlas
c62bad5fcd fix missing imports 2020-06-29 11:08:57 -07:00
Jake VanderPlas
afce718eb1 Add ability to specify individual test targets 2020-06-29 11:08:57 -07:00
Roy Frostig
6eea0ce17a detailed assertion in profiler test 2020-06-26 17:55:58 -07:00
Peter Hawkins
5116fd47aa
Add a heap profiler API and document it. (#3576) 2020-06-26 17:09:09 -04:00
Peter Hawkins
1298e9e8c4
Fix some test failures. (#2713) 2020-04-14 18:23:19 -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