| 3362 | } |
| 3363 | |
| 3364 | XlaOp AllReduce(const XlaOp operand, const XlaComputation& computation, |
| 3365 | absl::Span<const ReplicaGroup> replica_groups, |
| 3366 | const absl::optional<ChannelHandle>& channel_id, |
| 3367 | const absl::optional<Shape>& shape_with_layout) { |
| 3368 | return operand.builder()->AllReduce(operand, computation, replica_groups, |
| 3369 | channel_id, shape_with_layout); |
| 3370 | } |
| 3371 | |
| 3372 | XlaOp AllToAll(const XlaOp operand, int64 split_dimension, |
| 3373 | int64 concat_dimension, int64 split_count, |