57 Commits

Author SHA1 Message Date
Jake VanderPlas
ef45db7374 DOC: add stub for removed pjit tutorial 2023-02-07 08:44:56 -08:00
Jake VanderPlas
20b55a119e CI: update jupytext version 2023-01-23 14:42:03 -08:00
Yash Katariya
12f7cdeeae Delete the pjit 101 tutorial which is subsumed by https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html
PiperOrigin-RevId: 503581557
2023-01-20 20:35:47 -08:00
Ravin Kumar
2257e2075d Update doc landing page
Co-authored-by: 8bitmp3 <19637339+8bitmp3@users.noreply.github.com>
Co-authored-by: Jake VanderPlas <jakevdp@google.com>
2023-01-13 12:45:28 -08:00
Simon Butt
542b38a1c2 Updated jax.tree_leaves --> jax.tree_util.tree_leaves to remove deprecation notice in jax101-pytrees tutorial
Signed-off-by: Simon Butt <simonbutt123@gmail.com>
2022-12-03 21:22:25 +00:00
Igor Saprykin
be527b62d7 Make it clear that fun is a function rather than a noun.
PiperOrigin-RevId: 490370522
2022-11-22 17:02:21 -08:00
Peter Hawkins
71a384d25e Clarify in JAX Basics that JAX array creation is also an operation that requires accelerator dispatch and converting to a regular Python type is a blocking operation. 2022-10-26 13:38:17 -04:00
Jake VanderPlas
eeb9b5f1f6 pre-commit hook: update flake8, mypy, & jupytext 2022-08-15 15:32:45 -07:00
Ruffalo-sunghwan
d71581d2cb correct spelling on comment 2022-06-06 15:24:41 +09:00
Jake VanderPlas
29efc30c63 DOC: update myst-nb to v0.15.0 2022-05-09 14:32:48 -07:00
ngam
7132fbaca1 minor edits to clarify pjit docs 2022-05-09 16:14:45 +00:00
Jake VanderPlas
1134ff5871 CI: update jupytext to 0.13.8 2022-05-02 11:22:31 -07:00
Kehang Han
4717a6b6d2 updates maps.mesh() context manager in jax_101/pjit documentation
PiperOrigin-RevId: 445501580
2022-04-29 14:06:11 -07:00
Jake VanderPlas
c61a18b346 DOC: switch from tree_multimap to tree_map in docs 2022-04-01 14:52:16 -07:00
Jake VanderPlas
f4b64f48f4 doc: add examples of using partial with jit 2022-03-29 15:43:58 -07:00
Jake VanderPlas
b9b79bab31 maint: update pre-commit package versions & fix new mypy errors 2022-01-31 13:39:11 -08:00
Gerald Shen
4833fbb672 fix caching example in jit tutorial 2022-01-26 12:03:40 -05:00
Grigory Sapunov
504728d8b6 link directly to the documentation for the jnp.ndarray.at property 2021-12-29 12:29:16 +03:00
Grigory Sapunov
f93531b020 replace deprecated jax.ops.index_* functions with the new index update operators 2021-12-27 20:23:29 +03:00
Vlad Feinberg
cd333f0f5b Fix straight-through estimator example in docs (#9032) 2021-12-21 22:25:12 +00:00
Yash Katariya
821fcaa750 Make the pjit docs clear about who does local and global communication
PiperOrigin-RevId: 405421833
2021-10-25 09:37:15 -07:00
Jake VanderPlas
198cc5ee4f chore: update jupytext to v0.1.13 & re-sync notebooks 2021-10-05 14:30:16 -07:00
TinaZetongJia
bc01a3a5cc
removing ipynb in format 2021-09-09 17:03:35 -07:00
Tina Jia
ae8bbab798 add space 2021-09-09 00:10:48 +00:00
Tina Jia
4a36692792 added all imgs 2021-09-09 00:09:10 +00:00
Tina Jia
c0d9d76e8b added first img 2021-09-09 00:02:06 +00:00
Tina Jia
d18d24bb3d Merge branch 'pjit_tutorial' of https://github.com/TinaZetongJia/jax into pjit_tutorial 2021-09-08 23:50:43 +00:00
TinaZetongJia
849ce51328
delete images 2021-09-08 17:50:24 -06:00
Tina Jia
da186f4df7 delete pjit notebook 2021-09-08 23:38:05 +00:00
TinaZetongJia
851316f0ea
Uploaded smaller images 2021-09-07 18:50:35 -06:00
Jake VanderPlas
d3dd4e0ad3 don't use top-level headings 2021-09-04 06:41:16 -07:00
Jake VanderPlas
701ffe7286 sync md file 2021-09-04 06:29:14 -07:00
Jake VanderPlas
1edbf7636e Remove stderr blocks 2021-09-04 06:25:26 -07:00
Jake VanderPlas
a48cbbbe3b add markdown version 2021-09-03 19:01:16 -07:00
Tina Jia
9f1504bd72 added pjit tuorial colab 2021-09-03 17:49:58 -07:00
Tina Jia
24961d726f added pjit tuorial colab 2021-09-03 23:23:53 +00:00
elliotwaite
7392a57b75 DOC: many small fixes 2021-08-04 16:55:13 -07:00
Shawn Presser
76773798b2 Fix broken link in jax-101/06-parallelism 2021-06-29 17:02:45 -05: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
Jake VanderPlas
d96ec4b27a DOC: clean up JAX 101 notebooks and remove WIP warning 2021-04-01 16:27:06 -07:00
Skye Wanderman-Milne
624110e034 Reduce verbosity of jax-101 outline and fix pytree tutorial headers 2021-03-30 09:48:44 -07:00
Jake VanderPlas
f0b947bf6c JAX 101: add pytrees notebook 2021-03-22 10:21:02 -07:00
ebuehrle
2e90e72e90 Fix code typo 2021-03-15 17:20:47 +01:00
jax authors
65ee6041cb Merge pull request #6020 from jamestwebber:patch-1
PiperOrigin-RevId: 362360958
2021-03-11 13:11:21 -08:00
James Webber
cf4f44548c
sync markdown 2021-03-11 14:24:41 -05:00
Horace He
abdee43688
Update 04-advanced-autodiff.md 2021-03-11 05:56:52 -05:00
Horace He
e38d272357
Update 04-advanced-autodiff.ipynb 2021-03-11 05:55:48 -05:00
James Webber
3a99e58033
(hopefully) fix broken link 2021-03-10 20:09:54 -05:00
Jake VanderPlas
35933238d1 Initial commit of seventh JAX-101 notebook 2021-03-08 14:57:45 -08:00
jax authors
a68e8b3c76 Merge pull request #5963 from jakevdp:jax-101-part6
PiperOrigin-RevId: 361626184
2021-03-08 12:19:11 -08:00