add comment

openxla · Mar 19, 2024 · bbe5840 · bbe5840
1 parent 23508eb
commit bbe5840
Showing 1 changed file with 9 additions and 0 deletions.
diff --git 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();