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

88 Commits

Author SHA1 Message Date
Yash Katariya
0d5dae09ff Delete xmap and the jax.experimental.maps module. It's been 5 months since its deprecation (more than the standard 3 months deprecation period). 
PiperOrigin-RevId: 655614395
 2024-07-24 10:24:09 -07:00
Sharad Vikram
39ec5dacb4 [Pallas TPU] Add matrix multiplication tutorial 2024-07-16 18:12:19 -07:00
Sergei Lebedev
e80632e6fd Revived the workaround for not-expanding type aliases 
The version here only works for modules with
``from __future__ import annotations``, but we can safely add that import
to all modules now, since the minimal Python version JAX supports is 3.10.

The worakround was previously removed in #3485.
 2024-07-01 14:31:53 +01:00
Dan Foreman-Mackey
dda6430f7c Add register_custom_call_target to xla_client API docs. 
This function is (for better or worse) user facing for custom call
users. I think it's worth having this in the API docs.
 2024-06-27 14:40:36 -04:00
8bitmp3
15a3debee8 Add Distributed data loading doc 2024-05-16 23:12:13 +00:00
Jake VanderPlas
10ed827fe9 DOC: replace old tutorials with new content 2024-05-03 12:20:06 -07:00
Stephan Hoyer
c77370c492 Recursively pull out __wrapped__ in linkcode_resolve() 
This should actually fix the source lookup for `jax.nn.relu`, which
uses both `custom_jvp` and `jit` decorators.
 2024-05-03 09:10:00 -07:00
Stephan Hoyer
f6ee61273a Make linkcode_resolve() a bit more robust 
This handles three cases that came up when adopting this snippet (for
finding source code corresponding to API docs) [for neuralgcm](https://github.com/google-research/neuralgcm/pull/58):

1. documenting class method or attributes
2. documenting properties
3. documenting `jit` decorated methods

I'm not sure if case (1) or (2) comes up in the JAX docs, but case (3)
definitely does -- `jit` decorated functions like `jax.nn.relu`
[currently do not](https://jax.readthedocs.io/en/latest/_autosummary/jax.nn.relu.html)
have source code link.
 2024-05-02 17:57:02 -07:00
jax authors
17d0eb30fa Merge pull request #20961 from jakevdp:linkcode 
PiperOrigin-RevId: 629568292
 2024-04-30 16:56:14 -07:00
Jake VanderPlas
55e3be65b7 Remove dependency on sphinx_autodoc_typehints 2024-04-29 14:19:46 -07:00
Jake VanderPlas
74f1d8897c DOC: add manual documentation to jax.scipy.special functions. 
This lets us give more implementation-specific information, and
lets us avoid a needless dependency on scipy.
 2024-04-29 10:58:07 -07:00
Jake VanderPlas
8915ed8348 DOC: link docstrings to github source 2024-04-26 15:56:52 -07:00
Jake VanderPlas
8e34da70f8 DOC: add introduction to sharded computation 2024-04-16 10:39:18 -07:00
Sharad Vikram
2f8d5cebff [Pallas TPU] Add Pipelining and BlockSpecs documentation for Pallas TPU 2024-02-28 19:17:58 -08:00
8bitmp3
df3cc491ce
Update JAX docs copyright 2024-01-24 14:19:45 +00:00
Jake VanderPlas
a3abd178f0 DOC: run shard_map docs as part of the doc build 2024-01-22 18:07:01 -08:00
Matthew Johnson
8b219d5f7b [shard-map] add user tutorial 2024-01-18 15:30:13 -08:00
Jake VanderPlas
f29ec904f6 CI: fix doc build 2023-11-16 07:59:07 -08:00
Jake VanderPlas
a268863842 DOC: set navigation_with_keys=False 2023-10-25 10:08:13 -07:00
Jake VanderPlas
130a53f2a2 DOC: re-enable execution of thinking_in_jax.ipynb 2023-08-24 09:23:26 -07:00
Sharad Vikram
96e2d93f53 [Pallas] Add Pallas design doc 2023-08-03 16:02:07 -07:00
Sharad Vikram
7f1ef32ba3 Add initial documentation for Pallas 2023-08-03 12:30:19 -07:00
Jake VanderPlas
f1b0b6ac65 DOC: fix readthedocs for sphinx-book-theme=1.0 2023-03-01 10:11:31 -08:00
Matthew Johnson
141996ec11 add remat tutorial docs 2023-02-23 14:37:52 -08:00
Jake VanderPlas
a022a4e923 DOC: remove transformations.md 
It's currently unused, and the content duplicates what's in the README
 2023-02-07 12:32:11 -08:00
Jake VanderPlas
3ab0633d38 DOC: simplify jax-101 patterns in conf.py 2023-02-07 09:36:26 -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
Jake VanderPlas
7965c907a9 [typing] improve sphinx rendering of type aliases 2023-01-04 09:10:48 -08:00
Jake VanderPlas
406b5525a4 DOC: fix sphinx version for readthedocs 2022-12-29 15:37:54 -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
7b59ce2f89 DOC: pre-execute the quickstart notebook on GPU 2022-12-06 13:24:02 -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
yashkatariya
aca7e4ade2 jax.Array tutorial 2022-11-15 16:49:17 -08: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
Filippo Vicentini
52236adeed
Show link to GitHub repo in navbar 2022-09-05 13:41:49 +02:00
Peter Hawkins
71b29b1cc6 Create JAX Enhancement Proposals (JEPs). 
Migrate existing design documents to JEPs.
 2022-08-08 16:13:58 -04:00
Jake VanderPlas
d8d836bfb0 Fix execution issues in type_promotion doc 2022-07-11 14:16:11 -07:00
Ahti Kitsik
2764da5c71 DOC: fix docs generator to group notebook code output 2022-05-21 09:36:38 +02:00
Jake VanderPlas
c2f4a1a409 DOC: fix scipy intersphinx link 2022-05-20 15:17:26 -07:00
Jake VanderPlas
681240244c DOC: explicitly enable dollarmath extension 2022-05-12 10:37:34 -07:00
Jake VanderPlas
29efc30c63 DOC: update myst-nb to v0.15.0 2022-05-09 14:32:48 -07:00
tornikeo
6423741db0 Add copy button to code snippets in documentation 2022-02-22 13:49:48 +04:00
Jake VanderPlas
7381bbe8bb Add type promotion design doc 2022-02-14 14:16:42 -08: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
Peter Hawkins
efacc93088 Use the sphinx-book-theme for JAX documentation. 2022-02-04 09:00:07 -05:00
Roy Frostig
e659c54977 docs: remove application tutorials 2021-09-10 18:52:47 -07:00
TinaZetongJia
e7e8183a40 not exclude 08-pjit.md 2021-09-09 15:24:04 -07:00
George Necula
117d0d23ab Attempt to fix RTD build 
It seems that the failure is for transformations.md
 2021-07-18 19:35:14 +03:00
Peter Hawkins
b746d021c8 Disable xmap_tutorial to fix doc CI build. 2021-06-25 19:55:23 -04: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