Diff Coverage

    _OP_MAP['mean'] = dist.ReduceOp.AVG
else:
    # Fallback for older torch versions if necessary, though this might require manual division upstream
    # Assuming standard behavior where 'mean' implies native AVG support or upstream handling
    _OP_MAP['mean'] = dist.ReduceOp.SUM


def _ensure_contiguous(x):
    """Return a contiguous copy of *x* if not already contiguous."""

def _ensure_contiguous(x):
    """Return a contiguous copy of *x* if not already contiguous."""
    if torch.compiler.is_compiling():
        return x.contiguous()
    if not x.is_contiguous() or x.storage_offset() != 0:
        return x.contiguous()
    return x

        _TorchContiguousGrad.apply(tensor)
        for tensor in dist_func.all_gather(data, group=group)
    ]
    if rank_list is not None:
        group_ranks = dist.get_process_group_ranks(group)
        if tuple(rank_list) != tuple(group_ranks):
            rank_to_idx = {int(rank): idx for idx, rank in enumerate(group_ranks)}
            output = [output[rank_to_idx[int(rank)]] for rank in rank_list]
    return torch.cat(output, dim=concat_dim)


def differentiable_reduce_scatter(data, dev_num, axis, op, group):

    output_tensor = dist_func.reduce_scatter(output_tensor, input_tuple, op=reduce_op, group=group)

    # Keep manual handling for 'avg' string as it maps to SUM in _OP_MAP
    if op == 'avg':
        output_tensor = output_tensor / dev_num
    return output_tensor


def wait_async_tensor(tensor):