distributed.h

Go to the documentation of this file.

// Copyright © 2024 Apple Inc.
 
#pragma once
 
#include <memory>
 
#include "mlx/array.h"
 
namespace mlx::core::distributed {
 
// Forward declaration of the base group implementation.
namespace detail {
class GroupImpl;
};
 
/* Check if a communication backend is available */
bool is_available();

struct Group {
  Group(std::shared_ptr<detail::GroupImpl> group) : group_(std::move(group)) {}
 
  int rank() const;
  int size() const;

  Group split(int color, int key = -1) const;
 
  const std::shared_ptr<detail::GroupImpl>& raw_group() const {
    return group_;
  }
 
 private:
  std::shared_ptr<detail::GroupImpl> group_{nullptr};
};

Group init(bool strict = false, const std::string& bk = "any");
 
} // namespace mlx::core::distributed