Function all_reduce 

pub fn all_reduce<const D: usize>(
    input: Tensor<D>,
    op: ReduceOperation,
    device_ids: Vec<Device>,
) -> CollectiveTensor<D>

Performs an all_reduce operation on the input tensor.

Arguments

• input: The input tensor.
• op: The aggregation operation.
• device_ids: The list of all devices with which to all_reduce

Returns

A CollectiveTensor containing the handle of the result.