Adds send/recv ops in distributed (#1366)

mlx/distributed/ops.cpp

@@ -1,5 +1,7 @@
 // Copyright © 2024 Apple Inc.
 
+#include <sstream>
+
 #include "mlx/distributed/ops.h"
 #include "mlx/distributed/primitives.h"
 
@@ -57,4 +59,59 @@ array all_gather(
       {x});
 }
 
+array send(
+    const array& x,
+    int dst,
+    std::optional<Group> group_ /* = std::nullopt */,
+    StreamOrDevice s /* = {} */) {
+  auto group = to_group(group_);
+
+  if (group.size() == 1) {
+    throw std::invalid_argument("Cannot send to a singleton group");
+  }
+
+  if (dst < 0 || dst >= group.size()) {
+    std::ostringstream msg;
+    msg << "Invalid destination=" << dst << " for a group of size "
+        << group.size();
+    throw std::invalid_argument(msg.str());
+  }
+
+  return array(
+      {0}, int32, std::make_shared<Send>(to_stream(s), group, dst), {x});
+}
+
+array recv(
+    std::vector<int> shape,
+    Dtype dtype,
+    int src,
+    std::optional<Group> group_ /* = std::nullopt */,
+    StreamOrDevice s /* = {} */) {
+  auto group = to_group(group_);
+
+  if (group.size() == 1) {
+    throw std::invalid_argument("Cannot recv from a singleton group");
+  }
+
+  if (src < 0 || src >= group.size()) {
+    std::ostringstream msg;
+    msg << "Invalid source=" << src << " for a group of size " << group.size();
+    throw std::invalid_argument(msg.str());
+  }
+
+  return array(
+      std::move(shape),
+      std::move(dtype),
+      std::make_shared<Recv>(to_stream(s), group, src),
+      std::vector<array>{});
+}
+
+array recv_like(
+    const array& x,
+    int src,
+    std::optional<Group> group_ /* = std::nullopt */,
+    StreamOrDevice s /* = {} */) {
+  return recv(x.shape(), x.dtype(), src, group_, s);
+}
+
 } // namespace mlx::core::distributed