14 Commits

Author SHA1 Message Date
8bitmp3
0cf040c9a1 Add/update JAX Advanced Tutorials docs, ToC structure 2024-09-20 23:06:54 +00:00
jax authors
886aa944fa Merge pull request #23707 from jakevdp:stop-gradient-doc
PiperOrigin-RevId: 676876785
2024-09-20 09:48:08 -07: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
Jake VanderPlas
3c37d4f20e Improve documentation for jax.lax.stop_gradient 2024-09-17 15:58:14 -07:00
Frederik Wilde
d08b68996a
Update jax-primitives.md 2024-09-05 21:29:56 +02:00
Jake VanderPlas
09fd345de9 pre-commit: update hooks & pin using hashes 2024-08-27 15:23:13 -07:00
Jake VanderPlas
68be5b5085 CI: update ruff to v0.6.1 2024-08-27 14:54:11 -07:00
carlosgmartin
f812d0f28b Clarify meaning of tall and wide Jacobian matrices in autodiff docs. 2024-08-26 16:00:36 -07:00
Roy Frostig
371935cc10 update README and several docs to typed RNG keys 2024-08-11 08:09:47 -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
jax authors
fc1e1d4a65 Add freshness metablock to JAX OSS docs.
PiperOrigin-RevId: 645508135
2024-06-21 14:50:49 -07:00
Sergei Lebedev
01194bd2fb Clarified the type of the inputs to callback APIs
The callback APIs were migrated to use jax.Arrays for both inputs and outputs
in JAX 0.4.27.

PiperOrigin-RevId: 634473890
2024-05-16 11:29:09 -07:00
tempoxylophone
4bcad15111
Fix type in advanced-autodiff.md
removed `doc.g` text from the first section.
2024-05-14 11:52:41 -04:00
Jake VanderPlas
10ed827fe9 DOC: replace old tutorials with new content 2024-05-03 12:20:06 -07:00