jax.device_put#

jax.device_put(x, device=None)[source]#

Transfers x to device.

Parameters
  • x – An array, scalar, or (nested) standard Python container thereof.

  • device (Optional[Device]) – The (optional) Device to which x should be transferred. If given, then the result is committed to the device.

If the device parameter is None, then this operation behaves like the identity function if the operand is on any device already, otherwise it transfers the data to the default device, uncommitted.

For more details on data placement see the FAQ on data placement.

This function is always asynchronous, i.e. returns immediately.

Returns

A copy of x that resides on device.