2292557fd6
dist/include/ contains the MLX and MLX-C headers needed for CGo compilation. Without these, go-mlx cannot be used as a module dependency (headers not found in module cache). Libraries (dylib/metallib) are still gitignored — users build those locally via cmake. Co-Authored-By: Virgil <virgil@lethean.io>
56 lines
1.1 KiB
C++
56 lines
1.1 KiB
C++
// Copyright © 2024 Apple Inc.
|
|
|
|
#pragma once
|
|
|
|
#include <optional>
|
|
|
|
#include "mlx/distributed/distributed.h"
|
|
#include "mlx/utils.h"
|
|
|
|
namespace mlx::core::distributed {
|
|
|
|
array all_sum(
|
|
const array& x,
|
|
std::optional<Group> group = std::nullopt,
|
|
StreamOrDevice s = {});
|
|
|
|
array all_gather(
|
|
const array& x,
|
|
std::optional<Group> group = std::nullopt,
|
|
StreamOrDevice S = {});
|
|
|
|
array send(
|
|
const array& x,
|
|
int dst,
|
|
std::optional<Group> group = std::nullopt,
|
|
StreamOrDevice s = {});
|
|
|
|
array recv(
|
|
Shape shape,
|
|
Dtype dtype,
|
|
int src,
|
|
std::optional<Group> group = std::nullopt,
|
|
StreamOrDevice s = {});
|
|
|
|
array recv_like(
|
|
const array& x,
|
|
int src,
|
|
std::optional<Group> group = std::nullopt,
|
|
StreamOrDevice s = {});
|
|
|
|
array all_max(
|
|
const array& x,
|
|
std::optional<Group> group = std::nullopt,
|
|
StreamOrDevice s = {});
|
|
|
|
array all_min(
|
|
const array& x,
|
|
std::optional<Group> group = std::nullopt,
|
|
StreamOrDevice s = {});
|
|
|
|
array sum_scatter(
|
|
const array& x,
|
|
std::optional<Group> group = std::nullopt,
|
|
StreamOrDevice s = {});
|
|
|
|
} // namespace mlx::core::distributed
|