mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-15 19:36:06 +00:00
Code Issues Packages Projects Releases Wiki Activity
26,443 Commits 500 Branches 263 Tags
Commit Graph

103 Commits

Author SHA1 Message Date
Dougal
e8f43d1cef Explicit sharding docs 2025-03-14 16:33:30 -04:00
carlosgmartin
bc43b00d8f Add navigation breadcrumbs to docs. 2025-03-12 16:52:34 -04:00
Jake VanderPlas
380fb0a5a0 CI: skip execution of convolutions.ipynb 
This is failing on readthedocs with a dead kernel error
 2025-02-20 13:14:59 -08:00
Dougal
9145366f6f Part 1 of a new autodidax based on "stackless" 2025-02-12 15:23:06 -05:00
Dan Foreman-Mackey
21281f377e Remove the tutorial about legacy custom calls. 2025-02-06 21:04:30 -05:00
Axel Donath
dd7c203b81 Adjust conf.py to remove prompt from sphinx copybutton 2025-01-31 13:51:27 -05:00
Dan Foreman-Mackey
cb4d97aa1f Move jex.ffi to jax.ffi. 2024-12-29 13:06:19 +00:00
Jake VanderPlas
abf14323dc Adjust copyright notice. 
Previously we had been pulling-in NumPy and SciPy docs at runtime, but
after the work in #21461 this is no longer the case.
 2024-10-28 18:53:38 -07:00
8bitmp3
0cf040c9a1 Add/update JAX Advanced Tutorials docs, ToC structure 2024-09-20 23:06:54 +00:00
Michael Hudgins
d4d1518c3d Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax 
PiperOrigin-RevId: 676843138
 2024-09-20 07:52:33 -07:00
Justin Fu
4bce4f6452 [Pallas] Add block-sparse kernel tutorial 2024-09-19 12:23:03 -07:00
Jake VanderPlas
ab29fee763 Add array_api intersphinx & document jnp.permute_dims 2024-09-09 08:23:38 -07:00
Justin Fu
bdb03309b5
Merge branch 'main' into pallas_distr_docs 2024-08-08 17:36:36 -07:00
Justin Fu
dcd186f552 [Pallas] Add pallas distributed computation tutorial 2024-08-08 17:34:35 -07:00
Dan Foreman-Mackey
0b4800a193 Add ffi_call tutorial 
Building on #21925, this tutorial demonstrates the use of the FFI using
`ffi_call` with a simple example. I don't think this should cover all of
the most advanced use cases, but it should be sufficient for the most
common examples. I think it would be useful to eventually replace the
existing CUDA tutorial, but I'm not sure that it'll get there in the
first draft.

As an added benefit, this also runs a simple test (akin to
`docs/cuda_custom_call`) which actually executes using a tool chain that
open source users would use in practice.
 2024-08-01 15:36:32 -04:00
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