def broadcast_recv(shape, dtype, group_size, group_key, instance_key, communication_hint='auto'): """Receives a broadcasts tensor, across devices. Args: shape: Shape of the tensor to be received. dtype: Type of the tensor to be received. group_size: one plus the number of receiving tensors, i.e. the total number of devices participating. Each tensor must reside on a different device. group_key: an integer identifying the group of devices. instance_key: an integer identifying the participating group of Ops. communication_hint: preferred collective communication. The implementation may fall back to another mechanism. Options include `auto`, `ring`, and `nccl`. Returns: An Op implementing the broadcast receive. Raises: ValueError: if any of the input parameter constraints are not met. """ if group_size <= 1: raise ValueError( 'Parameter group_size to broadcast_send must be at least 2.') return gen_collective_ops.collective_bcast_recv( shape=shape, T=dtype, group_size=group_size, group_key=group_key, instance_key=instance_key, communication_hint=communication_hint.lower())
def broadcast_recv(shape, dtype, group_size, group_key, instance_key): """Receives a broadcasts tensor, across devices. Args: shape: Shape of the tensor to be received. dtype: Type of the tensor to be received. group_size: one plus the number of receiving tensors, i.e. the total number of devices participating. Each tensor must reside on a different device. group_key: an integer identifying the group of devices. instance_key: an integer identifying the participating group of Ops. Returns: An Op implementing the broadcast receive. Raises: ValueError: if any of the input parameter constraints are not met. """ if group_size <= 1: raise ValueError( 'Parameter group_size to broadcast_send must be at least 2.') return gen_collective_ops.collective_bcast_recv(shape=shape, T=dtype, group_size=group_size, group_key=group_key, instance_key=instance_key)
def broadcast_recv(shape, dtype, group_size, group_key, instance_key, communication_hint='auto', timeout=0): """Receives a broadcasts tensor, across devices. Args: shape: Shape of the tensor to be received. dtype: Type of the tensor to be received. group_size: one plus the number of receiving tensors, i.e. the total number of devices participating. Each tensor must reside on a different device. group_key: an integer identifying the group of devices. instance_key: an integer identifying the participating group of Ops. communication_hint: preferred collective communication. The implementation may fall back to another mechanism. Options include `auto`, `ring`, and `nccl`. timeout: If set to a non zero, set a completion timeout to detect staleness. If the timer goes off, a DeadlineExceededError is raised. The timeout value in seconds. This feature is experimental. Returns: An Op implementing the broadcast receive. Raises: ValueError: if any of the input parameter constraints are not met. """ if group_size <= 1: raise ValueError( 'Parameter `group_size` to broadcast_send must be at least 2. ' f'Received: {group_size}.') return gen_collective_ops.collective_bcast_recv( shape=shape, T=dtype, group_size=group_size, group_key=group_key, instance_key=instance_key, communication_hint=communication_hint.lower(), timeout_seconds=timeout)