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)