search

jax-ml / jax

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more
http://jax.readthedocs.io/
Apache License 2.0
30.56k stars 2.81k forks source link

An example for `ffi` input/output aliasing #24986

Open saeedmaleki opened 2 days ago

saeedmaleki commented 2 days ago

an example here to show how input/output are aliased.

dfm commented 1 day ago

Good suggestion! I've just opened a PR (https://github.com/jax-ml/jax/pull/25042) that adds such an example. I didn't add it to the tutorial, because we are using the examples/ffi directory in the JAX repo to demonstrate more advanced features for now. Please take a look at https://github.com/jax-ml/jax/pull/25042 (or, more precisely, this diff) to see if you have any more specific questions that aren't answered there.