def transfer_to_infeed(value):
"""Transfers the given value into the XLA infeed queue.
XLA's infeed queue is a single queue that feeds the "XLA virtual machine" with
a totally ordered stream of values. This is dequeued from XLA computations via
the Infeed() operation.
TODO(leary): this currently implicitly enqueues to device ordinal 0.
Args:
value: the value that the caller would like to enqueue into the XLA infeed
queue
"""
c_api.TransferToInfeedLocal(require_numpy_array_layout(value))
def transfer_to_infeed(value, replica_number=None):
"""Transfers the given value into the XLA infeed queue.
XLA's infeed queue is a single queue that feeds the "XLA virtual machine" with
a totally ordered stream of values. This is dequeued from XLA computations via
the Infeed() operation.
Args:
value: the value that the caller would like to enqueue into the XLA infeed
queue
replica_number: the replica number to infeed the value to -- if not
provided, then the default replica (trivially replica 0) is used.
"""
if replica_number is None:
c_api.TransferToInfeedLocal(require_numpy_array_layout(value))
else:
c_api.TransferToInfeedLocalReplica(
require_numpy_array_layout(value), replica_number)
