From bbe5840b8eb56a306a66ed03d701fd8976e01491 Mon Sep 17 00:00:00 2001
From: Clemens Giuliani <clemens@inailuig.it>
Date: Tue, 19 Mar 2024 17:30:13 +0100
Subject: [PATCH] add comment

---
 xla/pjrt/cpu/mpi_collectives.h | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/xla/pjrt/cpu/mpi_collectives.h b/xla/pjrt/cpu/mpi_collectives.h
index 6819af4b94309..e8934b4582df6 100644
--- a/xla/pjrt/cpu/mpi_collectives.h
+++ b/xla/pjrt/cpu/mpi_collectives.h
@@ -72,6 +72,15 @@ class MpiCollectivesCommunicator : public CollectivesCommunicator {
 class MpiCollectives : public CollectivesInterface {
  public:
 
+  /*
+  The user has to explicitly call Init() and Finalize() before and
+  after use.
+  For example, using the Python client, this can be achieved with:
+
+  collectives = xla_client._xla.make_mpi_collectives()
+  collectives.Init()
+  atexit.register(collectives.Finalize)
+  */
   void Init();
   void Finalize();