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

186 Commits

Author SHA1 Message Date
Richard Levasseur
f891cbf64b Load Python rules from rules_python 
PiperOrigin-RevId: 559789250
 2023-08-24 10:22:57 -07:00
George Necula
26f091e446 [callback] Disable stream_executor tests. 
PiperOrigin-RevId: 559252832
 2023-08-22 16:15:00 -07:00
Peter Hawkins
a259df0d76 Move compiler APIs out of dispatch.py and xla_bridge.py into a new jax._src.compiler module. 
Refactoring only, no user-visible changes intended.

PiperOrigin-RevId: 557116160
 2023-08-15 06:39:46 -07:00
Jake VanderPlas
ad8e719b82 Add jnp.ufunc and jnp.frompyfunc 2023-08-10 14:58:18 -07:00
Peter Hawkins
afd56c15d9 Move jax.jaxpr_util to jax._src.jaxpr_util, and split it into a separate build target. 
Change jaxpr_util_test to be a py_test(), since there's no point testing it on every hardware configuration.

PiperOrigin-RevId: 554861284
 2023-08-08 10:09:09 -07:00
Adam Paszke
0228bf7d3c Fix MSAN errors in cache_key_test 
The device_assignment array was never initialized, causing MSAN errors.
Replacing it with np.arange fixes the issue.

PiperOrigin-RevId: 553469463
 2023-08-03 07:28:32 -07:00
jax authors
9e7502ce60 Disable MSAN testing for cache-key unit tests. 
This is an existing issue with the compilation cache tests.
The refactoring of the cache key generation part into a separate
file requires tagging the refactored tests also.

PiperOrigin-RevId: 551972670
 2023-07-28 16:02:12 -07:00
Skye Wanderman-Milne
8b58e38ec5 Add jax_debug_log_modules config option. 
This can be used to enable debug logging for specific files
(e.g. `JAX_DEBUG_LOG_MODULES="jax._src.xla_bridge,jax._src.dispatch"`)
or all jax (`JAX_DEBUG_LOG_MODULES="jax"`).

Example output:
```
$ JAX_DEBUG_LOG_MODULES=jax python3 -c "import jax; jax.numpy.add(1,1)"
DEBUG:2023-06-07 00:27:57,399:jax._src.xla_bridge:352: No jax_plugins namespace packages available
DEBUG:2023-06-07 00:27:57,488:jax._src.path:29: etils.epath found. Using etils.epath for file I/O.
DEBUG:2023-06-07 00:27:57,663:jax._src.dispatch:272: Finished tracing + transforming fn for pjit in 0.0005719661712646484 sec
DEBUG:2023-06-07 00:27:57,664:jax._src.xla_bridge:590: Initializing backend 'tpu'
DEBUG:2023-06-07 00:28:00,502:jax._src.xla_bridge:602: Backend 'tpu' initialized
DEBUG:2023-06-07 00:28:00,502:jax._src.xla_bridge:590: Initializing backend 'cpu'
DEBUG:2023-06-07 00:28:00,542:jax._src.xla_bridge:602: Backend 'cpu' initialized
DEBUG:2023-06-07 00:28:00,544:jax._src.interpreters.pxla:1890: Compiling fn for with global shapes and types [ShapedArray(int32[], weak_type=True), ShapedArray(int32[], weak_type=True)]. Argument mapping: (GSPMDSharding({replicated}), GSPMDSharding({replicated})).
DEBUG:2023-06-07 00:28:00,547:jax._src.dispatch:272: Finished jaxpr to MLIR module conversion jit(fn) in 0.0023522377014160156 sec
DEBUG:2023-06-07 00:28:00,547:jax._src.xla_bridge:140: get_compile_options: num_replicas=1 num_partitions=1 device_assignment=[[TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0)]]
DEBUG:2023-06-07 00:28:00,571:jax._src.dispatch:272: Finished XLA compilation of jit(fn) in 0.023587703704833984 sec
```
 2023-07-28 18:11:12 +00:00
jax authors
3b28d4e180 Refactor Jax compilation cache key generation. 
This is in preparation for introducing a more robust key-generation
algorithm.

This refactoring does not introduce any change in behavior.

Testing: refactored unit tests and test workload.
PiperOrigin-RevId: 551744892
 2023-07-27 23:01:00 -07:00
Sharad Vikram
3d556b7a19 Add Mosaic to Jaxlib and expose bindings in jax.experimental.mosaic 
PiperOrigin-RevId: 549801858
 2023-07-20 18:28:51 -07:00
Tao Wang
b7686f41aa Enable passing fdo_profile in compiler_options in pxla.py 
PiperOrigin-RevId: 549109629
 2023-07-18 14:18:28 -07:00
Peter Hawkins
3fcf72af32 Disable tests that are flaky in CI. 
PiperOrigin-RevId: 547464866
 2023-07-12 05:13:38 -07:00
jax authors
2fa6a9c9bf Allow other backends to run the array_test.py test. 
PiperOrigin-RevId: 547191886
 2023-07-11 08:05:25 -07:00
Chris Flesher
5be17ed90c Added scipy.spatial.transform Rotation and Slerp classes 2023-06-08 07:51:32 -05:00
Peter Hawkins
32026ad18b Disable random_test_with_custom_prng on CPU under msan. 
This test flakily times out in CI.

PiperOrigin-RevId: 535293997
 2023-05-25 10:10:01 -07:00
Peter Hawkins
1d20d2f301 Increase sharding of host_callback_test on TPU to fix CI flakiness. 
PiperOrigin-RevId: 533451822
 2023-05-19 07:44:53 -07:00
Peter Hawkins
e6628e2e72 Disable tests that time out in CI. 
PiperOrigin-RevId: 532792740
 2023-05-17 08:16:07 -07:00
Peter Hawkins
9471bb3045 Disable sparsify_test on CPU under tsan. 
Under tsan this test times out in CI.

PiperOrigin-RevId: 531210930
 2023-05-11 08:33:35 -07:00
Peter Hawkins
e8c735125c Disable more tests that are flaky in CI. 
PiperOrigin-RevId: 529724306
 2023-05-05 08:33:33 -07:00
pizzud
40d730be49 aot_test: Stop forcing XLA to assume a certain number of devices. 
Test cases are still frequently skipped due to lack of CompileOptions
support, but the skip/run behavior does not seem to meaningfully change
compared to a clean checkout. This was verified by inserting an exception
in place of unittest.SkipTest.

PiperOrigin-RevId: 529437419
 2023-05-04 09:53:26 -07:00
Peter Hawkins
09fce87f54 Increase sharding of or disable some flaky CI tests. 
PiperOrigin-RevId: 529405705
 2023-05-04 07:41:56 -07:00
Peter Hawkins
57e62ca03c Reenable scipy_stats_test in CI. 
Disable testTruncNormPdf on CPU, which is failing after an LLVM update.

PiperOrigin-RevId: 528884880
 2023-05-02 14:11:08 -07:00
Skye Wanderman-Milne
70cac773f7 Exclude scipy_fft_test from msan as well as t/asan. 
PiperOrigin-RevId: 528562775
 2023-05-01 13:42:24 -07:00
Skye Wanderman-Milne
fa68c1f882 Bump up lax_test TPU sharding to avoid asan timeouts 
PiperOrigin-RevId: 528559870
 2023-05-01 13:31:22 -07:00
Skye Wanderman-Milne
c662fd216d Disable tsan CI for random_test_with_custom_prng to avoid timeouts. 
asan is already disabled, and the comment and "cpu" case indicates
that tsan should already have been disabled as well.

PiperOrigin-RevId: 528000458
 2023-04-28 15:26:46 -07:00
Skye Wanderman-Milne
67d80c21cb Increase sharding count on nn_test and svd_test to avoid ASAN timeouts. 
PiperOrigin-RevId: 527387005
 2023-04-26 15:11:29 -07:00
Jake VanderPlas
1c7f8efce6 Add test framework for module attribute 2023-04-21 13:20:16 -07:00
jax authors
db2cbd4ae8 Merge pull request #15665 from hawkinsp:sourceinfo 
PiperOrigin-RevId: 525581713
 2023-04-19 16:30:23 -07:00
Peter Hawkins
3bb7386149 [JAX] Improve handling of metadata in compilation cache. 
Metadata, in particular code location information is present in the HLO generated by JAX. The compilation cache uses the serialized HLO as a cache key, which begs the question: should code location information be part of that key? Simply changing the line number on which a function appears shouldn't necessarily cause a cache miss.

There are pros and cons: the main advantage of excluding metadata is that we will get more cache hits, and the main disadvantage is that debug information and profiling data in the HLO might become confusing, since it may refer to a different program entirely, or to a version of a program that does not correspond to the current state of the source tree. We argue that saving compilation time is the more important concern.

This change adds a tiny MLIR pass that strips Locations from a StableHLO module, and applies it in the compilation cache if metadata stripping is enabled.

PiperOrigin-RevId: 525534901
 2023-04-19 13:27:04 -07:00
Peter Hawkins
a3b262c379 Use the traceback of the call site when assigning a source location to an inlined function. 
Improves but does not completely fix https://github.com/google/jax/issues/15663 . The non-inlined case still has similar problems.
 2023-04-19 13:56:53 -04:00
Peter Hawkins
017548c40b Move implementation of compilation cache out of jax/experimental and into jax/_src. 
Use a Protocol instead of an abstract base class for the CacheInterface since it allows us to use one fewer file.

No functional change intended.

PiperOrigin-RevId: 524855263
 2023-04-17 08:35:53 -07:00
Jake VanderPlas
8c8f50f688 Fix tolerance and shard_count for experimental_rnn_test 
This should fix the current GPU test timeout.

PiperOrigin-RevId: 522167894
 2023-04-05 15:19:19 -07:00
jax authors
3c1f3abba2 Merge pull request #15149 from sharadmv:runstate 
PiperOrigin-RevId: 521809360
 2023-04-04 10:56:25 -07:00
Parker Schuh
c2b15a1eb8 Break out aot_test from array_test (for serialization and other aot APIs). 
PiperOrigin-RevId: 521568985
 2023-04-03 14:47:53 -07:00
Peter Hawkins
47177e1417 Split more targets out the main JAX Bazel target. 
Namely:
* abstract_arrays
* ad_util
* api_util
* interpreters/partial_eval
* lax_reference
PiperOrigin-RevId: 520618715
 2023-03-30 06:12:45 -07:00
Skye Wanderman-Milne
4cb3b011a0 Remove PJRT C API bypass. 
Now that all functionality needed by frameworks is implemented, let's
remove the possibility of not noticing missing functionality due to
the bypass.

PiperOrigin-RevId: 519018438
 2023-03-23 18:39:14 -07:00
Jieying Luo
b403c2a083 [PJRT C API] Add parsing PJRT client create options from json file. 
PiperOrigin-RevId: 518418760
 2023-03-21 16:57:34 -07:00
Yash Katariya
207cc10058 Error if jax_array or jax_jit_pjit_api_merge is set to False. 
PiperOrigin-RevId: 517485597
 2023-03-17 12:57:57 -07:00
Peter Hawkins
dea7450e4e Remove references to jax.config.jax_array, which is always True at head. 
PiperOrigin-RevId: 516970232
 2023-03-15 17:09:11 -07:00
Yash Katariya
88584290aa Remove GDA tests from JAX since GDA is deprecated. There are jax.Array tests for all the corresponding GDA tests 
PiperOrigin-RevId: 516881635
 2023-03-15 11:34:57 -07:00
Peter Hawkins
a32a7ff903 Move _src/tree_util.py into a separate Bazel target. 
Fix a type error in api.py revealed by the split.

PiperOrigin-RevId: 515745227
 2023-03-10 14:51:52 -08:00
pizzud
04def0b6ab lazy_loader_module: Move to new internal_test_util directory. 
Now we no longer need to mess with sys.path in lazy_loader_test.

PiperOrigin-RevId: 515674188
 2023-03-10 10:29:33 -08:00
Peter Hawkins
01b00c4821 Increase sharding of shard_map test on CPU. 
This test is timing out in CI with sanitizers enabled (asan/tsan).

PiperOrigin-RevId: 515369731
 2023-03-09 10:13:26 -08:00
jax authors
59bf2061c4 Merge pull request #14565 from pizzud:deprecation-module 
PiperOrigin-RevId: 515172435
 2023-03-08 16:23:53 -08:00
jax authors
9c4db8c962 Merge pull request #14633 from mattjj:shmap-test-vmap 
PiperOrigin-RevId: 515117185
 2023-03-08 12:56:54 -08:00
pizzud
22cbf95e07 lax_vmap_test: Extend timeout so that the TPU variant can run in ASAN. 
Unfortunately we can't conditionally change the timeout, as size and timeout
are both non-configurable even if jax_test supported setting the size.

PiperOrigin-RevId: 514745247
 2023-03-07 08:49:42 -08:00
jax authors
00f1abe401 Disable 2 failing jax tests. 
PiperOrigin-RevId: 514515343
 2023-03-06 13:50:40 -08:00
pizzud
ef28dcf091 lax_scipy_test: Split into three targets, take 2. 
The goal is to ensure that all shards fit into a medium timeout in sanitizer
configurations.

Running 256 entry vectors in spectral_dac is too slow, so let's replace that
with a smaller vector that isn't a power of 2. Avoiding a power of 2 requires
us to widen the tolerance a bit due to vectorization changes.

While here, specify deps a little more precisely as well.

PiperOrigin-RevId: 514440062
 2023-03-06 09:53:23 -08:00
pizzud
0292f5d0a6 lax_scipy_test: Revert split into three targets. 
Somehow the spectral_dac functionality is flaky on its own when run on CPU.

PiperOrigin-RevId: 512195860
 2023-02-24 16:56:40 -08:00
pizzud
09afbac6ff lax_scipy_test: Split into three so that each target is small enough to fit within a medium timeout. 
The spectral_dac tests are also shrunk because running the full suite on 256-entry vectors is too slow.

This allows them to run in ASAN in more situations.

While here, specify deps a little more precisely as well.

PiperOrigin-RevId: 511829646
 2023-02-23 10:51:58 -08:00