jax
d41078fb - Properly pack and unpack int4 arrays on CPU in PJRT.

Commit
2 years ago
Properly pack and unpack int4 arrays on CPU in PJRT. Transferring an array from host to device on CPU sometimes does a zero-copy implementation where no memory is actually moved. This is now never done with int4, since int4 arrays are stored in packed format on device and an unpacked format on host. Similarly, transferring an array from device to host on CPU used to always use a zero-copy implementation, but now it will unpack and copy for int4 arrays. PiperOrigin-RevId: 578692796
Author
Committer
Parents
Loading