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

7 Commits

Author SHA1 Message Date
Skye Wanderman-Milne
120125f3dd Make pytest-xdist work on TPU and update Cloud TPU CI. 
This change also marks multiaccelerator test files in a way pytest can
understand (if pytest is installed).

By running single-device tests on a single TPU chip, running the test
suite goes from 1hr 45m to 35m (both timings are running slow tests).

I tried using bazel at first, which already supported parallel
execution across TPU cores, but somehow it still takes 2h 20m! I'm not
sure why it's so slow. It appears that bazel creates many new test
processes over time, vs. pytest reuses the number of processes
initially specified, and starting and stopping the TPU runtime takes a
few seconds so that may be adding up. It also appears that
single-process bazel is slower than single-process pytest, which I
haven't looked into yet.
 2022-11-18 22:05:13 +00: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
Yash Katariya
e6bdb00d31 Skip remote_transfer_test because Array does not have the xla_shape method since its deprecated. 
PiperOrigin-RevId: 474913967
 2022-09-16 15:26:12 -07:00
Peter Hawkins
0150d15cb2 Increase minimum jaxlib version to 0.3.7. 
Drop backwards compatibility with older jaxlib versions.
 2022-04-18 08:09:50 -04:00
Peter Hawkins
7ffdac0746 Skip remote_transfer_test on cloud TPU. 
The necessary API isn't yet implemented for Cloud TPU.

PiperOrigin-RevId: 442058546
 2022-04-15 11:24:52 -07:00
Peter Hawkins
b3a62cd3f2 Disable remote_transfer_test on GPU. It currently crashes. 
PiperOrigin-RevId: 441762941
 2022-04-14 08:06:27 -07:00
Jonathan Heek
df20bd2de5 Expose CopyToRemoteDevice and MakeCrossHostReceiveBuffer in Python bindings. 
PiperOrigin-RevId: 441746248
 2022-04-14 06:40:48 -07:00