distributed_impl.h

Go to the documentation of this file.

// Copyright © 2024 Apple Inc.
 
#pragma once
 
#include "mlx/distributed/distributed.h"
 
namespace mlx::core::distributed::detail {
 
/* Return the communication stream. */
Stream communication_stream();
 
/* Perform an all reduce sum operation */
void all_sum(Group group, const array& input, array& output);
 
/* Perform an all gather operation */
void all_gather(Group group, const array& input, array& output);
 
void send(Group group, const array& input, int dst);
 
void recv(Group group, array& out, int src);
 
} // namespace mlx::core::distributed::detail