mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-24 19:26:05 +00:00
Code Issues Packages Projects Releases Wiki Activity
14,805 Commits 508 Branches 263 Tags
Commit Graph

193 Commits

Author SHA1 Message Date
Jake VanderPlas
c4ec2996af Sharp bits: mention alternatives to lax.cond 2023-02-02 13:19:26 -08:00
Jake VanderPlas
fd71794633 DOC: organize content in advanced guide 2023-01-26 08:20:53 -08:00
Jake VanderPlas
20b55a119e CI: update jupytext version 2023-01-23 14:42:03 -08:00
Jake VanderPlas
9e355a6606 Sharp Bits: add section on Dynamic shapes 2023-01-19 11:37:03 -08:00
Jake VanderPlas
b2141229e3 Sharp bits: raise exceptions directly 2023-01-17 14:13:38 -08:00
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