This allows JAX to communicate on-device arrays with other libraries, such as PyTorch and CuPy. PiperOrigin-RevId: 290845329 Change-Id: Idd99d81533159bc2ad0c5177b69ac7f30315cb1a