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

39 Commits

Author SHA1 Message Date
jax authors
48a55a6d71 Add a profiler test for gpu run 
PiperOrigin-RevId: 732247572
 2025-02-28 13:45:46 -08:00
Peter Hawkins
c61b2f6b81 Make JAX test suite pass (at least most of the time) with multiple threads enabled. 
Add a new jtu.thread_unsafe_test_class() decorator to tag entire `TestCase` classes as thread-hostile.

PiperOrigin-RevId: 714037277
 2025-01-10 06:58:46 -08:00
Adam Paszke
7c1bea6b2a Increase the timeout in profiler tests, but use threading.Event to exit as soon as possible 
The previous timeout ended up being too small for Cloud TPUs. I bumped it to 30s, which seems
to be sufficient, but I also made sure to break out of the loop as soon as the profile completes.

PiperOrigin-RevId: 707523546
 2024-12-18 06:08:32 -08:00
Ram Rachum
799c79ca94 Support pathlib.Path input to start_trace 2024-07-19 23:41:02 +03:00
Jake VanderPlas
f090074d86 Avoid 'from jax import config' imports 
In some environments this appears to import the config module rather than
the config object.
 2024-04-11 13:23:27 -07:00
Jieying Luo
ca1844dd60 [PJRT C API] Bump the minimum support libtpu version and clean up version check. 
PiperOrigin-RevId: 610508159
 2024-02-26 13:26:01 -08:00
jax authors
9b27d43e70 Import submodules from jax._src explicitly, instead of relying on import side-effects. It will lead to the missing x-refs in code search according to go/pywald-sawmill-analysis. 
PiperOrigin-RevId: 604788105
 2024-02-06 15:47:16 -08:00
Peter Hawkins
30a0136813 Increase minimum jaxlib version to 0.4.19. 
0.4.19 has xla_extension version 207 and mlir_api_version 54.

PiperOrigin-RevId: 583412447
 2023-11-17 09:38:31 -08:00
Jieying Luo
21d9e50641 Disable some more profiler test for older plugins. 
PiperOrigin-RevId: 582071616
 2023-11-13 13:47:50 -08:00
Jieying Luo
11236dbe34 Disable profiler test for older plugins. 
PiperOrigin-RevId: 581391435
 2023-11-10 15:58:20 -08:00
Tao Wang
93fbf623a4 Fix testProfilerGetFDOProfile. 
PiperOrigin-RevId: 573936890
 2023-10-16 14:52:44 -07:00
Tao Wang
c568110cc8 Set up an API to top trace and fdo profile in memory. 
PiperOrigin-RevId: 573276173
 2023-10-13 11:34:25 -07:00
Peter Hawkins
1885c4933c Add a new internal test utility test_device_matches() and use it instead of equality tests on device_under_test(). 
This change prepares for allowing more flexible tag matching. For example, we may want to write "gpu" in a test and have it match both "cuda" and "rocm" devices, which we cannot do under the current API but can easily do under this design.

Replace uses of device_under_test() in a context that performs an equality test with a call to test_device_matches().
Replace uses of if_device_under_test() with test_device_matches() and delete if_device_under_test().

PiperOrigin-RevId: 568923117
 2023-09-27 12:10:43 -07:00
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