jax

13b1bca9 - prevent device_put from being staged into jaxprs