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

1138 Commits

Author SHA1 Message Date
Roy Frostig
e0ddd43933 add custom_vjp upgrade flag to jax module and fix doc rendering for upgrade flags 2022-03-28 20:36:35 -07:00
Hyeontaek Lim
87671f4d85 Add a documentation for the JAX transfer guard 2022-03-25 19:50:12 +00:00
jax authors
f57e78e240 Merge pull request #9766 from ayaka14732:patch-1 
PiperOrigin-RevId: 437297267
 2022-03-25 11:39:41 -07:00
jax authors
440cd7b9df Merge pull request #9906 from jakevdp:gamma-log-space 
PiperOrigin-RevId: 436255847
 2022-03-21 11:11:50 -07:00
Jake VanderPlas
32c718db11 quickstart: fix outdated TPU reference 2022-03-21 09:41:10 -07:00
Jake VanderPlas
69969ef803 add random.loggamma and improve dirichlet & beta implementation 2022-03-21 08:33:11 -07:00
Matthew Johnson
43036e1316 improvements to autodidax reduce_sum 
* generalize reduce_sum to handle multiple axes
* add reduce_sum transpose rule

also fix bug in AD jaxpr formation related to deduplicating consts
 2022-03-15 20:22:55 -07:00
Jake VanderPlas
09105b2d58 Update outdated jax.numpy module docs 2022-03-15 11:02:59 -07:00
Jake VanderPlas
c66f5dda60 DOC: add missing linalg functionality to docs 2022-03-15 09:55:59 -07:00
Peter Hawkins
80aec7b25f Documentation improvements. 2022-03-08 09:37:33 -05:00
Peter Hawkins
d3d666d081 Document jax.nn.initializers. 2022-03-07 17:26:04 -05:00
jax authors
f87e6c3a74 Merge pull request #9622 from hawkinsp:docs4 
PiperOrigin-RevId: 432440609
 2022-03-04 07:29:29 -08:00
Ayaka Mikazuki
ee2377b47d
[Doc] Add jaxlib build instructions for TPU 
Currently the documentation only contains build instructions for CUDA support, but does not mention TPU. This pull request add build instructions for TPU alongside CUDA.
 2022-03-04 20:22:59 +08:00
jax authors
2e158c1362 Merge pull request #9764 from jakevdp:doc-doctest 
PiperOrigin-RevId: 432278100
 2022-03-03 14:02:15 -08:00
Jake VanderPlas
f1ccb35d8c DOC: add doctest info to developer docs 2022-03-03 13:11:00 -08:00
yashkatariya
f0df7b7348 Fix the display of GDA docs 2022-03-03 12:19:04 -08:00
jax authors
d369501417 Merge pull request #9546 from pnkraemer:jet-docs 
PiperOrigin-RevId: 432073995
 2022-03-02 18:04:40 -08:00
jax authors
d9f82f7b9b [JAX] Move experimental.ann.approx_*_k into lax. 
Updated docs, tests and the example code snippets.

PiperOrigin-RevId: 431781401
 2022-03-01 14:46:33 -08:00
Jake VanderPlas
ed2550999f implement jnp.copy 2022-03-01 11:56:36 -08:00
Peter Hawkins
2b51ab5bc1 Add an API compatibility policy. 2022-03-01 14:28:22 -05:00
Jake VanderPlas
74e7fdfd4c DOC: update sharp bits info about fori_loop differentiability 2022-02-25 09:15:55 -08:00
jax authors
d2834b64df Merge pull request #9658 from tornikeo:docs-copybutton 
PiperOrigin-RevId: 430705653
 2022-02-24 08:12:41 -08:00
Peter Hawkins
f51a05a889 Remove jax.ops.index... functions. 
These functions have been deprecated and have issued a DeprecationWarning since jax 0.2.22 in October 2021.
 2022-02-24 09:36:28 -05:00
tornikeo
6423741db0 Add copy button to code snippets in documentation 2022-02-22 13:49:48 +04:00
Peter Hawkins
3e5ecfe363 Add jax.distributed and jax.dlpack to the docs. 
Reorder the doc modules into something closer to alphabetical order.

Add missing functions from jax.scipy.linalg and jax.scipy.signal to the docs.
 2022-02-17 16:10:07 -05:00
jax authors
95c486a5a3 Merge pull request #9585 from jakevdp:typos 
PiperOrigin-RevId: 428871943
 2022-02-15 14:16:43 -08:00
Jake VanderPlas
e82f232ea9 Fix nomenclature in type promotion doc 2022-02-15 13:26:06 -08:00
Jake VanderPlas
de360833e0 type promotion design doc: minor typos 2022-02-15 12:17:18 -08:00
Jake VanderPlas
7381bbe8bb Add type promotion design doc 2022-02-14 14:16:42 -08:00
Nicholas Krämer
faf57d813b Wrote docstring for jet 2022-02-12 08:14:57 +01:00
Yash Katariya
1ad3551ec9 Release jax and jaxlib 0.3.0 as per the new release process. 
PiperOrigin-RevId: 427809845
 2022-02-10 11:59:13 -08:00
Lena Martens
0042edb5f4 Checkify: rename some symbols and add some docstrings. 2022-02-08 17:40:04 +00:00
Peter Hawkins
8be057de1f Introduce a new jax/jaxlib versioning scheme. 
Adds a design note that describes the scheme and how the jax and jaxlib versions
are related.
 2022-02-07 17:59:42 -05:00
Peter Hawkins
465b593293 Update scipy intersphinx inventory for SciPy 1.8.0. 
According to https://github.com/scipy/scipy/issues/14267 the SciPy docs seems to have moved.
 2022-02-07 16:19:46 -05:00
Jake VanderPlas
ea8817b329 DOC: move experimental APIs to their own pages 2022-02-04 14:40:34 -08:00
Jake VanderPlas
fc10438b4f DOC: move functions in jax.html to their own pages 2022-02-04 14:40:34 -08:00
jax authors
248572c3e8 Merge pull request #9446 from hawkinsp:docsj 
PiperOrigin-RevId: 426408109
 2022-02-04 08:23:21 -08:00
jax authors
2a61403845 Merge pull request #9440 from hawkinsp:booktheme 
PiperOrigin-RevId: 426407795
 2022-02-04 08:18:12 -08:00
Peter Hawkins
efacc93088 Use the sphinx-book-theme for JAX documentation. 2022-02-04 09:00:07 -05:00
Peter Hawkins
a43c82a3b7 Add -j auto to the suggested doc build instructions. 
Despite being marked as experimental, parallelism appears to work fine for JAX doc builds, and speeds up builds significantly.
 2022-02-04 08:57:14 -05:00
jax authors
45d96c490e Merge pull request #4671 from romanngg:conv_local 
PiperOrigin-RevId: 426282505
 2022-02-03 18:03:33 -08:00
Peter Hawkins
df55ea5204 Move design notes into docs/, and render them as part of the documentation. 2022-02-02 14:29:03 -05:00
Jake VanderPlas
b9b79bab31 maint: update pre-commit package versions & fix new mypy errors 2022-01-31 13:39:11 -08:00
jax authors
d66daa9039 Merge pull request #9312 from GeraldCSC:jit-tutorial-update 
PiperOrigin-RevId: 424875694
 2022-01-28 08:27:39 -08:00
Jake VanderPlas
928087ada0 DOC: add info about repeated indices to jax.ops docs 2022-01-27 10:22:25 -08:00
Gerald Shen
4833fbb672 fix caching example in jit tutorial 2022-01-26 12:03:40 -05:00
Jake VanderPlas
6c38ec9a05 developer doc: more info on pre-commit 2022-01-25 10:36:03 -08:00
Jake VanderPlas
42a2e66fbe DOC: pin docutils==0.16 to restore bullets in lists 2022-01-24 12:02:39 -08:00
jax authors
7f91763303 Merge pull request #9287 from google:gda 
PiperOrigin-RevId: 423866601
 2022-01-24 11:26:37 -08:00
George Necula
83b818d45c Add more documentation for buffer donation 
Fixes: #9237
 2022-01-24 09:33:08 +01:00