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
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
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
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
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
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
Jake Vanderplas
05904faf0f
Change onp/np to np/jnp in docs & notebooks ( #3760 )
2020-07-15 13:17:38 -07:00
Jin Dong
a23a6c2a06
remove pip install
2020-04-09 18:16:38 -04:00
Matthew Johnson
9afa2c6b69
fix broken link to trax, fixes #1974
2020-01-10 20:44:24 -08:00
Skye Wanderman-Milne
1f88c831bb
Add Cloud TPU Colabs landing page.
2019-12-13 15:48:32 -08:00
