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

45 Commits

Author SHA1 Message Date
Michael Hudgins
d4d1518c3d Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax 
PiperOrigin-RevId: 676843138
 2024-09-20 07:52:33 -07:00
Jake VanderPlas
68be5b5085 CI: update ruff to v0.6.1 2024-08-27 14:54:11 -07:00
Roy Frostig
371935cc10 update README and several docs to typed RNG keys 2024-08-11 08:09:47 -07:00
Benjamin Bastian
c980dc4e2a
Update outdated comment Pmap_Cookbook 
Colab now has TPU VMs, so this notebook works on Colab!
 2024-04-19 12:44:15 -07:00
Benjamin Bastian
afc87785a6
Remove outdated section in cloud_tpu_colabs README 
Colab now has TPU VMs, so this section is now out of date.
 2024-04-18 15:48:13 -07:00
Jan Hrček
4da56dcdd7 Fix duplicate word occurrences 2023-12-19 06:15:30 +01:00
8bitmp3
ed2fef8a29 Add Kaggle integration for JAX pmap cookbook 2023-11-06 20:37:52 +00:00
Peter Hawkins
2c32660a8f Replace references to DeviceArray with Array. 
A number of stale references are lurking in our documentation.
 2023-08-18 17:46:00 -04:00
Peter Hawkins
6ed66ada0f Delete remote TPU support. 
TPU VMs are the only supported way to use TPUs as of JAX 0.4.0.

PiperOrigin-RevId: 519211267
 2023-03-24 12:33:33 -07:00
Peter Hawkins
f66f6ec98a [JAX] Move jax._src.lib.xla_bridge to jax._src.xla_bridge. 
Limit jax._src.lib to shims around jaxlib and nothing else.

The goal of this change is to avoid a dependency cycle between the rest of jax and jax._src.lib in a Bazel build. This allows the types for jax._src.lib to be inferred by pytype in isolation without referring to the rest of JAX.

PiperOrigin-RevId: 512922397
 2023-02-28 07:01:57 -08: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
Jake VanderPlas
c61a18b346 DOC: switch from tree_multimap to tree_map in docs 2022-04-01 14:52:16 -07:00
Peter Hawkins
9ea55468ab [JAX] Update users of jax.ops.index... functions, which are deprecated. 
* replace uses of `jax.ops.index[...]` with `jax.numpy.index_exp[...]`, which is a standard NumPy function that does the same thing.
* remove some redundant uses of `jax.ops.index[...]`, where the expression is passed directly to an indexed accessor function like `.at[...]`.
* update some remaining users of `jax.ops.index_update(x, jax.ops.index[idx], y)` to use the `x.at[idx].set(y)` APIs.

PiperOrigin-RevId: 406162068
 2021-10-28 09:54:26 -07:00
Peter Hawkins
11127add2b [JAX] Replace users of deprecated jax.lib.xla_bridge functions with more public names. 
Change in preparation for removing some function exports from jax.lib.xla_bridge.

PiperOrigin-RevId: 404849899
 2021-10-21 12:07:39 -07:00
Peter Hawkins
2c2f4033cc Move contents of jax.lib to jax._src.lib. 
Add shim libraries for functions exported from jax.lib that other code seems to use in practice.

PiperOrigin-RevId: 398471863
 2021-09-23 06:33:55 -07:00
Jake VanderPlas
d0e21072ec DOC: add brief comments to example in README 2021-09-14 11:47:19 -07:00
George Necula
6a48c60a72 Rename master to main in embedded links. 
Tried to avoid the change on external links to repos that
have not yet renamed master.
 2021-06-18 10:00:01 +03:00
Skye Wanderman-Milne
8a65ff8c25 Update cloud_tpu_colabs README to reflect Cloud TPU VM announcement. 2021-06-02 16:55:52 -07:00
Jake VanderPlas
24e920e3f5 change colab.sandbox->colab.research 2021-02-25 10:29:43 -08:00
George Necula
d6247defd6
Update Pmap_Cookbook.ipynb 2021-01-31 22:19:13 +02:00
George Necula
580a8f54cf
Update Pmap_Cookbook.ipynb 
I am doing this mostly to trigger another run of CI; there was a problem with our CI runs and this PR appears to be failing tests.
 2021-01-31 16:23:39 +02:00
mat kelcey
4e48023538
update out of date comment in Pmap_Cookbook.ipynb 
update comment regarding `pmap` not support `in_axes` because it does now.
 2021-01-31 09:07:55 +11:00
James Bradbury
1e2c5551db
fix dropped import 2021-01-13 11:07:50 -08:00
Skye Wanderman-Milne
85796cc7e3 Fix colab_tpu.setup_tpu import in example Cloud TPU notebooks. 2020-12-14 16:23:20 -08:00
Skye Wanderman-Milne
e802bf9b63 Merge pull request #5173 from jakevdp:conv-elem-type-cleanup 
PiperOrigin-RevId: 347085811
 2020-12-11 15:27:40 -08:00
Skye Wanderman-Milne
9f4015d41a Add NeurIPS 2020 demo notebook. 2020-12-07 20:03:29 +00:00
Skye Wanderman-Milne
a6b0709873 Demo notebook placeholder 2020-12-04 15:46:51 -08:00
Jake Vanderplas
05904faf0f
Change onp/np to np/jnp in docs & notebooks (#3760) 2020-07-15 13:17:38 -07:00
Yash Katariya
aaf0c4e36b
Install scikit-image and add try;except around from google.colab import files import statement as google.colab is not available locally and in kokoro (#3626) 2020-07-01 11:41:24 -07:00
Yash Katariya
236a5881bf
Add try; except to cells that fail and print out the error message so that the notebook can execute end-to-end (#3624) 2020-07-01 08:56:10 -07:00
Matthew Johnson
e9ce700b06
update tpu readme not to pin jax / jaxlib versions (#3467) 2020-06-16 15:48:01 -07:00
Jin Dong
a23a6c2a06 remove pip install 2020-04-09 18:16:38 -04:00
Skye Wanderman-Milne
74faef18bf
Rename "NeurIPS 2019 JAX demo" to "JAX demo", as it's been presented elsewhere. (#2372) 2020-03-06 10:32:15 -08:00
Ron Weiss
a473e5b6bb
Reference deduped precision issue. (#2319) 2020-02-27 10:20:01 -05:00
Skye Wanderman-Milne
f6d483739a
Shrink nested pmap image in demo notebook. (#2293) 2020-02-22 23:10:13 -08:00
Skye Wanderman-Milne
7885b1d034
Add nested pmap image to demo notebook. (#2292) 2020-02-22 23:04:34 -08:00
Skye Wanderman-Milne
5e77789afe Update Cloud TPU email address to jax-tpu@googlegroups.com 2020-02-11 13:42:11 -08:00
Matthew Johnson
9afa2c6b69 fix broken link to trax, fixes #1974 2020-01-10 20:44:24 -08:00
Skye Wanderman-Milne
a6da7e8bf9
Link to "Reporting issues" at top of Cloud TPU README. 2019-12-16 09:58:30 -08:00
Matthew Johnson
c952ed4f5f
update readme to clarify Cloud TPU preview status 2019-12-16 09:32:55 -08:00
Matthew Johnson
6ba39f763d tweak demo 2019-12-14 10:33:37 -08:00
Matthew Johnson
5c800367d1 mention cloud tpus in readme 2019-12-14 10:32:04 -08:00
Skye Wanderman-Milne
1efe648d35 Add demo colab. 2019-12-13 16:17:58 -08:00
Skye Wanderman-Milne
1f88c831bb
Add Cloud TPU Colabs landing page. 2019-12-13 15:48:32 -08:00