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

33 Commits

Author SHA1 Message Date
jax authors
01819257f6 Merge pull request #11701 from sharadmv:state 
PiperOrigin-RevId: 464658336
 2022-08-01 17:05:43 -07:00
Yash Katariya
9a5af235da Delete sharded_jit 
PiperOrigin-RevId: 464081692
 2022-07-29 08:19:52 -07:00
Peter Hawkins
9e6254e058 Increase shard counts for TPU tests in an attempt to fix CI timeouts under asan. 
PiperOrigin-RevId: 463830139
 2022-07-28 07:14:36 -07:00
George Necula
07fcf79324 jax.mask and jax.shapecheck are being deprecated 
Issue: #11557
PiperOrigin-RevId: 462315754
 2022-07-21 00:09:31 -07:00
Kuangyuan Chen
c0ec3b33e6 Introduce jax.experimental.clear_backends to delete all JAX runtime backends. 
In cases like unit tests, users may want to clean up all the backends along with the resources used in the end of the test, and reinitialize them in the next test.

PiperOrigin-RevId: 462239974
 2022-07-20 15:10:27 -07:00
Adam Paszke
117da44712 Internal change 
PiperOrigin-RevId: 462110048
 2022-07-20 04:31:21 -07:00
George Necula
777c129dfb [dynamic-shapes] Split dynamic_api_test.py 
PiperOrigin-RevId: 461109288
 2022-07-14 20:18:53 -07:00
jax authors
3eff9d11d2 Internal change 
PiperOrigin-RevId: 460434859
 2022-07-12 05:21:20 -07:00
Peter Hawkins
64e0b5d801 Increase bazel sharding of GPU tests. 
Reduces the maximum time for some test shards to avoid flaky timeouts.
 2022-07-11 14:19:43 +00:00
Sharad Vikram
b666f665ec Rollback of HCB GPU custom call due to internal failures 
PiperOrigin-RevId: 460079787
 2022-07-10 13:05:27 -07:00
jax authors
66ab792fc0 Merge pull request #11383 from YouJiacheng:Enable-HCB-customCall-implementation-on-GPU 
PiperOrigin-RevId: 459872063
 2022-07-08 18:23:16 -07:00
YouJiacheng
7c707832aa Enable CustomCall implementation on GPU 2022-07-09 02:29:08 +08:00
Peter Hawkins
5a7bedca37 Increase shard_count for sparse_test_gpu to 20. 
1918d39765 updated the wrong test!

This test is close to the timeout in the GPU CI and flakes sometimes.

PiperOrigin-RevId: 459762867
 2022-07-08 08:30:26 -07:00
Peter Hawkins
1918d39765 Increase number of shards for GPU sparse_test to 20. 2022-07-07 21:14:25 -04:00
Sharad Vikram
6274b9ed39 Enable Python callbacks on TFRT TPU backend 
PiperOrigin-RevId: 459415455
 2022-07-06 20:52:50 -07:00
Peter Hawkins
95e79332c0 Add JAX_TEST_TARGETS and JAX_EXCLUDE_TEST_TARGETS environment variables to assist with skipping tests under Bazel. 
Add "multiaccelerator" test tags to mark tests that would meaningfully run with more than one accelerator (e.g., GPU).

PiperOrigin-RevId: 459320212
 2022-07-06 12:51:43 -07:00
Peter Hawkins
1fc9afd03a Add support for running JAX tests under Bazel. 
This is an alternative method for running the tests that some users may prefer: pytest is and will remain fully supported.

To use this, one creates a .bazelrc by running the existing `build.py` script, and then one can run the tests by running:
```
bazel test -c opt //tests/...
```

Issue #7323

PiperOrigin-RevId: 458551208
 2022-07-01 15:07:22 -07:00
Roy Frostig
28828772b8 remove unuzed bazel build rules, including bazel test definitions 2018-12-16 10:42:50 -08:00
Roy Frostig
a97f820961 update examples/BUILD 2018-12-15 15:54:31 -08:00
Peter Hawkins
3561b432c2 Add Cholesky, QR, and Triangular solve implementations. 
* Adds lax.{cholesky,triangular_solve,qr}. Adds a JVP for Cholesky.
* Adds a transpose rule for add_p, needed by the Cholesky JVP.
* Adds np.linalg.{cholesky,qr,dot,matmul,trace}.
* Adds scipy.linalg.{cholesky,qr,solve_triangular,tril,triu}.

Pair programmed with mattjj.
 2018-12-13 13:03:08 -05:00
Roy Frostig
cfc37eb0be add a test of the resnet50 example 2018-12-12 16:09:20 -08:00
Roy Frostig
9c69014318 Only build-depend on libjax optionally (via bazel CLI flag) in the generated bazel test targets. 2018-12-12 18:01:44 -05:00
Dougal Maclaurin
c3374a9d5f Added build rule for generated_fun_test (formerly quickish_check) 2018-12-06 17:04:00 -05:00
Dougal Maclaurin
29113dd606 Made tests runnable with bazel 2018-12-06 17:00:47 -05:00
Roy Frostig
f5b051b431 Double gpu test shards 
PiperOrigin-RevId: 223647959
 2018-12-02 11:50:46 -08:00
Roy Frostig
20878c76f4 source sync 
PiperOrigin-RevId: 223530503
 2018-12-02 11:50:39 -08:00
Peter Hawkins
6361b784a8 source sync 
PiperOrigin-RevId: 222456068
 2018-11-21 20:22:56 -08:00
Matthew Johnson
25fb9b421d source sync 
PiperOrigin-RevId: 222170151
 2018-11-21 20:22:33 -08:00
Roy Frostig
a3619ca89d source sync 
PiperOrigin-RevId: 222153576
 2018-11-21 20:22:30 -08:00
Matthew Johnson
50038c07c8 fix build file issues 2018-11-19 20:18:31 -08:00
Roy Frostig
99f98f8e8c source sync 2018-11-19 13:50:57 -08:00
Roy Frostig
da2d53ad33 source sync 2018-11-19 13:29:47 -08:00
Matthew Johnson
9ae0f3a610 split BUILD file, move up license files 2018-11-18 15:43:09 -08:00