dlpack.callback now supports mutable_results=

google / jax

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more

http://jax.readthedocs.io/

Apache License 2.0

28.23k stars 2.58k forks source link

dlpack.callback now supports mutable_results= #21025

Open copybara-service[bot] opened 2 weeks ago

copybara-service[bot] commented 2 weeks ago

dlpack.callback now supports mutable_results=

The new flag passes DLPack tensors for both arguments and results into the callback, allowing to do zero-copy interop with other tensor programming libraries.

See #20701 for motivation.