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

358 Commits

Author SHA1 Message Date
Peter Hawkins
c125aa7f19 Add Python 3.12.0-rc.3 to the Windows wheel releases. 2023-09-19 21:34:51 -04:00
Yash Katariya
cf3fc55da3 Put JAXLIB_RELEASE env var in the correct step (build wheel) for windows release 
PiperOrigin-RevId: 566461339
 2023-09-18 17:39:31 -07:00
Yash Katariya
4fb0b4bace Set JAXLIB_RELEASE env var to True for windows release scripts 
PiperOrigin-RevId: 566442791
 2023-09-18 16:17:48 -07:00
Jake VanderPlas
f8031fed5b CI: pin actions/checkout to v4 2023-09-05 09:26:15 -07:00
Skye Wanderman-Milne
719cae5849 Small Cloud TPU CI debugging improvements 
1. Use `strings` to print the libtpu build date, which is useful if
   the libtpu version is stale and causes an API version mismatch
   error.

2. Set the default shell to add bash -x command, so all commands are
   printed inline. This makes lining up the output with the command
   easier.
 2023-08-14 10:37:26 -07:00
Skye Wanderman-Milne
0f30685dac Remove StreamExecutor-based TPU runtime from Cloud TPU CI 
The old StreamExecutor-based backend is no longer supported as of
3e50fea29e
 2023-08-09 10:05:46 -07:00
Sharad Vikram
d872812a35 [Pallas] Upstream pallas to JAX 
PiperOrigin-RevId: 552963029
 2023-08-01 16:43:13 -07:00
Peter Hawkins
319ab98980 Apply pyupgrade --py39-plus. 
Notable changes:
* use PEP 585 type names
* use PEP 604 type union syntax where `from __future__ import annotations` is present.
* use f-strings in more places.
* remove redundant arguments to open().
 2023-07-21 14:49:44 -04: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
Skye Wanderman-Milne
6c909760d5 Cloud TPU CI: make sure we update test deps and upgrade protobuf version 
`profiler_test.py:ProfilerTest.test_remote_profiler` fails with the
protobuf upgrade. However, I was seeing mysterious hangs without this,
and in general I think we should be testing with up-to-date deps given
that we don't pin. I'm gonna continue working on getting the Cloud TPU
CI green.
 2023-07-19 16:17:47 -07:00
Peter Hawkins
816ba91263 Use lower-case PEP 585 names for types. 
Issue https://github.com/google/jax/issues/16537

PiperOrigin-RevId: 542969282
 2023-06-23 15:12:14 -07:00
Peter Hawkins
487b640acf Jax 0.4.13 release. 2023-06-22 14:59:36 -07:00
jax authors
961d918883 Merge pull request #16333 from skye:profiler_test 
PiperOrigin-RevId: 542069606
 2023-06-20 15:40:10 -07: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
Peter Hawkins
741f0d03c3 Rename Github actions so they have more consistent naming. 2023-06-20 16:26:12 -04:00
Peter Hawkins
34d9f5a9ae Add a CI presubmit that renders the documentation. 2023-06-20 09:29:25 -04:00
Peter Hawkins
3efe4f94ee Don't allow the use of stale cache entries in Github actions CI. 
The use of restore-keys allows the pip cache to match stale cache entries. This allows for partial cache reuse, but means the cache also grows without bound.
At the time of writing, the pip cache is 4.7GB, which is much larger than just downloading the wheels we need (20MB)! So it's probably a net positive not to reuse caches even if it requires some redundant downloads.
 2023-06-16 17:10:57 -04:00
Peter Hawkins
a05ffab786 Fix xla_bridge_test.py test failures. 
We are splitting the plugins in the enviroment variable using os.pathsep; we should make sure to use that as the separator in the test.
 2023-06-16 10:17:07 -04:00
Peter Hawkins
990eb2ce68 Enable colored output in pytest under github actions. 2023-06-15 18:57:12 -04:00
Peter Hawkins
758f93723a Tag the wheels generated by the Windows CI as dev versions. 2023-06-14 11:50:33 -04:00
Jake VanderPlas
bc0cbd9b3b CI nightly: update nightly wheel location 2023-06-14 02:30:56 -07:00
Peter Hawkins
7025800b3f Windows CI: don't fail-fast. 
We expect some of the tests to fail at the moment, and we'd like all of the builds to run to completion even if one of them fails.
 2023-06-13 20:09:51 -04:00
Peter Hawkins
2c71349502 Window CI: fix syntax error. 
Try escaping backslashes in bazel override path.

Upload wheels even if tests fail, since they may be useful for debugging.
 2023-06-13 18:19:40 -04:00
Peter Hawkins
7c1fb60d11 Windows CI fixes: 
* try single-quoting the bazel override path to avoid \ substitution.
* use a second github checkout action to checkout the XLA repository.
 2023-06-13 17:46:00 -04:00
Peter Hawkins
00a086a07d Windows CI: fix bazel repository override command. 2023-06-13 17:12:30 -04:00
Peter Hawkins
48f32a6b70 Fix Windows CI build. 2023-06-13 17:00:01 -04:00
Peter Hawkins
0a70d42fb9 Windows CI improvements: 
* upload wheel artifacts to the github workflow.
* use XLA from github head in CI builds.
* update release builds to match the CI build configuration.
 2023-06-13 16:46:20 -04:00
Peter Hawkins
e47c475fe4 Windows CI: write wheel into usual location. 
Should fix installation failure.
 2023-06-13 14:23:59 -04:00
Peter Hawkins
800463d4c7 Fix command line arguments to build.py. 2023-06-13 13:14:36 -04:00
Peter Hawkins
6b9df28568 Try another name for the Windows CI runner. 2023-06-13 13:01:14 -04:00
Peter Hawkins
38fefe6035 Update OS in windows CI to use 16 core variant. 2023-06-13 12:02:10 -04:00
Peter Hawkins
0ebbdc97e3 Use larger runners for Windows CI. 
Simplify build command line: we don't need to work around low-memory low-CPU configurations.
 2023-06-13 11:28:42 -04:00
Peter Hawkins
bb966c7768 Don't include the --subcommands on the Bazel command line for Windows CI. 
This makes the output very verbose.
 2023-06-12 16:53:20 -04:00
Peter Hawkins
72c3d0a193 Update Windows CI. 
For the wheel_win_x64.yml builds:
* Don't trigger a full wheel build on push. The branch name was wrong anyway ("master" instead of "main"), but likely this action is too expensive to for each PR.
* Disable builds on release for the moment, until we have validated this approach.

Add a new nightly and manually-triggerable Windows build, for a specific Python version.
 2023-06-12 16:27:21 -04:00
jax authors
cf223d3d87 Merge pull request #15009 from mlxd:add_ghactions_winx64 
PiperOrigin-RevId: 539737845
 2023-06-12 13:02:29 -07:00
Skye Wanderman-Milne
d480e16152 Add maxfail=20 to some pytest commands, take 2 
Somehow I messed up https://github.com/google/jax/pull/16195 and it didn't fully stick.
 2023-06-06 11:24:49 -07:00
Jake VanderPlas
1d7dced221 CI: set maxfail=20 for nightly tests 2023-06-02 03:43:29 -07:00
jax authors
727c121169 Merge pull request #16188 from nouiz:ci_jestimator 
PiperOrigin-RevId: 536810121
 2023-05-31 13:41:29 -07:00
Skye Wanderman-Milne
131d28ba0f Use default Python version on Cloud TPU CI 2023-05-31 18:04:39 +00:00
Frederic Bastien
cdced240f5 WAR the bug in t5x dependency. It currently need the dev version of jestimator. 2023-05-30 12:55:21 -07:00
Yash Katariya
009af38697 Remove gda_serialization from JAX. It's replacement is array_serialization 
PiperOrigin-RevId: 531525875
 2023-05-12 09:23:05 -07:00
Yash Katariya
954cda9ce1 Move lint_and_typecheck and documentation job to the ubuntu-latest image since we don't need a large machine for it 
PiperOrigin-RevId: 530734120
 2023-05-09 15:47:22 -07:00
jax authors
47f5c225a9 Merge pull request #15789 from nouiz:ci 
PiperOrigin-RevId: 529121113
 2023-05-03 10:02:01 -07:00
Jake VanderPlas
2fa2f82274 CI: run lint_and_typecheck under newest Python version 2023-05-02 12:09:53 -07:00
Frederic Bastien
8f1711a431 Install as all the other. 2023-05-02 08:36:08 -07:00
Jake VanderPlas
8a4eadfec5 Change Python 3.8 CI runs to Python 3.9 
Followup for #15805

PiperOrigin-RevId: 528495690
 2023-05-01 09:47:51 -07:00
Frederic Bastien
ce8a31f9e3 Try to install correctly the JAX nightly. 2023-04-28 13:31:14 -07:00
Yash Katariya
dcd9127598 Make CI nightly faster which tests numpy and scipy nightly builds. 
PiperOrigin-RevId: 525817202
 2023-04-20 12:22:32 -07:00
Frederic Bastien
bc0c25c4b5 pytest* just got removed. The CI don't need them anymore, so remove that requirement. 2023-04-19 08:26:04 -07:00
Frederic Bastien
30f30dd31d Print the error 2023-04-17 07:49:08 -07:00