188 Commits

Author SHA1 Message Date
8bitmp3
2bf6d3dace Update jax.xmap notebook title (Named axes and easy-to-revise parallelism with xmap) 2023-01-07 00:05:46 +00:00
yashkatariya
008a6918e0 Remove experimental endpoints and update to point to 0.4.1 2022-12-13 11:53:56 -08:00
Matthew Johnson
1185c895ca in jax.Array notebook, polish beginning and tweak title and some wording 2022-12-10 22:16:54 -08:00
Jake VanderPlas
df02d7035e DOC: add example of pure_callback with custom_jvp 2022-12-09 12:43:04 -08:00
Jake VanderPlas
7b59ce2f89 DOC: pre-execute the quickstart notebook on GPU 2022-12-06 13:24:02 -08:00
yashkatariya
70d50814b1 Add cross-linking for the migration guide and the parallelism with JAX
tutorial

Co-authored-by: Skye Wanderman-Milne <skyewm@google.com>
2022-12-01 14:42:59 -08:00
Skye Wanderman-Milne
51db1cfd0e [docs] Rename "JAX in Parallelism" files so the URL matches the title. 2022-12-01 19:53:31 +00:00
Roy Frostig
b6fd3ff9d7 describe partitionable RNG mode in parallelism doc
Co-authored-by: Yash Katariya <yashkatariya@google.com>
2022-11-29 14:31:06 -08:00
Roy Frostig
fcce6b102c remove cotangent negation in custom VJP example
This was originally intended to show that we can change the VJP by
customizing it, but the algebraic incorrectness is confusing.
2022-11-22 17:55:22 -08:00
yashkatariya
aca7e4ade2 jax.Array tutorial 2022-11-15 16:49:17 -08:00
Adam Paszke
6e43ce363e Remove a TODO from the xmap tutorial
xeinsum is already powerful enough to support the example.
2022-10-26 15:44:06 +00:00
Serge Durand
27360b9988 Fix book links 2022-10-03 13:59:10 +02: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
Jake VanderPlas
be65694ac6 docs: avoid deprecated matplotlib axis creation 2022-09-19 12:55:18 -07:00
Jake VanderPlas
eeb9b5f1f6 pre-commit hook: update flake8, mypy, & jupytext 2022-08-15 15:32:45 -07:00
Peter Hawkins
71b29b1cc6 Create JAX Enhancement Proposals (JEPs).
Migrate existing design documents to JEPs.
2022-08-08 16:13:58 -04:00
Matthew Johnson
7ed4fa550e in sharp bits notebook, add back iaml blog post link 2022-08-04 08:29:26 -07:00
Roy Frostig
cd77debfa7 fix broken link in "NNs with TFDS" notebook 2022-07-28 19:04:08 -07:00
Ian McKenzie
0cc2ada432 Fix broken links for moved design_notes folder 2022-06-24 12:18:11 -07:00
Qiao Zhang
be71989af6 Remove broken image link.
PiperOrigin-RevId: 456628284
2022-06-22 15:54:36 -07:00
Jake VanderPlas
ed152953ef DOC: fix output repr in thinking_in_jax 2022-06-21 12:54:50 -07:00
fehiepsi
3ca0d3f149 Rename mesh into Mesh in xmap tutorial and doc 2022-06-03 17:11:09 +07:00
Jake VanderPlas
8a3bfe0300 DOC: add references for haskell-style signatures 2022-05-13 12:35:28 -07:00
Sharad Vikram
aca9dc6c6f Update custom interpreter tutorial 2022-05-11 11:11:47 -07:00
Jake VanderPlas
29efc30c63 DOC: update myst-nb to v0.15.0 2022-05-09 14:32:48 -07:00
Jake VanderPlas
dd71e2920b DOC: remove mentions of units 2022-05-04 16:08:44 -07:00
Jake VanderPlas
1134ff5871 CI: update jupytext to 0.13.8 2022-05-02 11:22:31 -07:00
Jake VanderPlas
72f14a5d64 Sharp bits: add miscellaneous divergences from numpy 2022-03-31 12:15:06 -07:00
Jake VanderPlas
32c718db11 quickstart: fix outdated TPU reference 2022-03-21 09:41:10 -07:00
Jake VanderPlas
74e7fdfd4c DOC: update sharp bits info about fori_loop differentiability 2022-02-25 09:15:55 -08:00
Jake VanderPlas
b9b79bab31 maint: update pre-commit package versions & fix new mypy errors 2022-01-31 13:39:11 -08:00
Matthew Johnson
726f60f6bc remove cpu platform setting in quickstart
fixes #9244
2022-01-19 10:24:09 -08:00
Mike McCoy
2e5ab11652 Resolves issue 8744 2022-01-12 21:10:45 +00:00
che-shr-cat
78977d6f5a fix broken links and update texts in thinking_in_jax.ipynb 2022-01-10 16:19:57 +03:00
Jake VanderPlas
4a72e57ce0 DOC: add FAQ section on JAX vs. Numpy performance 2021-11-24 12:04:02 -08:00
Peter Hawkins
58199b4b9a Delete the XLA in Python notebook.
Its tests are failing, and it describes a non-public API that we are phasing out.
2021-11-18 09:45:06 -05:00
Peter Hawkins
e783cbcb72 Port remaining translation rules inside JAX to new style.
PiperOrigin-RevId: 404288551
2021-10-19 09:48:37 -07:00
Basil
6fac5d36a1 Remove unused variable 2021-10-13 14:38:33 +01:00
Peter Hawkins
104a46594b Add DeprecationWarnings to jax.ops.index_... operators.
Remove uses of index_... in Common Gotchas notebook.
2021-10-05 20:47:22 -04:00
Jake VanderPlas
198cc5ee4f chore: update jupytext to v0.1.13 & re-sync notebooks 2021-10-05 14:30:16 -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
jax authors
394ef618d5 Merge pull request #7868 from PWhiddy:main
PiperOrigin-RevId: 397197685
2021-09-16 16:43:35 -07:00
PWhiddy
af22aae1e0 fix html formatting of markdown 2021-09-15 16:13:02 -04:00
Qiao Zhang
aec40b48d3 Update docs for nan checker and tfds example. 2021-09-15 11:54:48 -07:00
Roy Frostig
e659c54977 docs: remove application tutorials 2021-09-10 18:52:47 -07:00
Jake VanderPlas
9959d087c1 maml.ipynb: fix sync issue 2021-08-24 09:11:40 -07:00
Xiaohu Zhu
db0efd917f
Update maml.md
typo fixed.
2021-08-21 15:50:36 +08:00
elliotwaite
7392a57b75 DOC: many small fixes 2021-08-04 16:55:13 -07:00
Shawn Presser
7b879c811d Fix typo in xmap_tutorial 2021-06-29 17:09:57 -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