12 Commits

Author SHA1 Message Date
Dan Foreman-Mackey
729418094e Add a discussion of sharding to the FFI tutorial. 2025-01-09 11:24:40 -05:00
Dan Foreman-Mackey
cb4d97aa1f Move jex.ffi to jax.ffi. 2024-12-29 13:06:19 +00:00
Jérome Eertmans
d912034cb5
fix(docs): typos in macro name
chore(docs): sync .md file
2024-11-19 16:42:19 +01:00
Dan Foreman-Mackey
f757054267 Update some outdated syntax in FFI tutorial. 2024-11-12 08:34:24 -08:00
Dan Foreman-Mackey
61701af4a2 Rename vmap methods for callbacks. 2024-10-21 15:03:04 -04:00
Dan Foreman-Mackey
0b651f0f45 Make ffi_call return a callable 2024-10-21 12:16:57 -04:00
Dan Foreman-Mackey
1d27d420ac Deprecate the vectorized argument to pure_callback and ffi_call. 2024-10-02 11:33:51 -04: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
Dan Foreman-Mackey
bcbc0962bb Add the FFI functions and tutorial to the changelog.
Although we soft launched the FFI with v0.4.31, it would be nice to
include an update in the changelog to help with visibility.
2024-09-06 12:30:28 -04:00
Roy Frostig
3c223cd253 docs: tidy up titles and headings
This shortens some titles and makes them more consistent. It also
removes "JAX" from several titles ("in JAX", "for JAX", "JAX's",
etc.). Since these are JAX docs, that ought to be clear from context.
2024-08-13 11:53:57 -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