defall_reduce(t, group_size, group_key, instance_key, merge_op, final_op, subdiv_offsets=(0,)) """Reducestensorscollectively, acrossdevices. Args: t: the tensor to be reduced. group_size: the total number of tensors to be collectively reduced. Each 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. merge_op: string naming the binary Op to be applied to compute each partial reduction. final_op: string naming the unary Op to be applied to each fully reduced value. Can be 'Id'for no operation. subdiv_offsets: a list of integer offsets into the tensor at which each independent subdivision should begin. Use [0] if no subdivision should be done.
Returns: An Op implementing the distributed reduction.
Raises: ValueError: ifany of the input parameter constraints are not met. """ def broadcast_send(t, shape, dtype, group_size, group_key, instance_key) """Broadcasts one tensor to a group of others, across devices.
Args: t: the tensor to be sent. shape: the shape of the tensor being sent, which must agree with t. dtype: the type of the tensor being sent, which must agree with t. 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 distributed broadcast send.
Raises: ValueError: ifany of the input parameter constraints are not met.
Note that the shape and dtype arguments appear redundant since they should be obtainable from t. The are two reasons for including them. First, the shape andtype of tensors passed via broadcast must be known ahead of time in their most specific form so that the receive side can allocate memory for the operation and shape/type inference can carry forward from there. Including the same declarations on the send side clarifies a commitment already made. Secondly, having nearly identical use syntax for send and receive sides may simplify tool-driven generation of broadcast. """ 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: ifany of the input parameter constraints are not met. """