jax

2abb37c2 - move a _device_put_raw under broadcast impl

move a _device_put_raw under broadcast impl Before this change, we had to interact with the device to construct an array of zeros, even if we were staging everything out (e.g. with jax.xla_computation and omnistaging).