jax
5583f652 - Add transparent copy of buffers in the C++ jax.jit.

Commit
5 years ago
Add transparent copy of buffers in the C++ jax.jit. If jax.jit is not committed to a device, an error is raise when several sticky buffers are on different devices. When jax.jit is committed to a device, we will move the data unconditionally. This should mimic the Python logic. PiperOrigin-RevId: 333130052
References
Author
Committer
Parents
Loading