Fix JAX extension build with NVTE_UB_WITH_MPI=1

Signed-off-by: Gaetan Lepage <gaetan@glepage.com>

Author: GaetanLepage

build_tools/jax.py

@@ -3,6 +3,7 @@
 # See LICENSE for license information.
 
 """JAX related extensions."""
+
 import os
 from pathlib import Path
 from packaging import version
@@ -100,6 +101,14 @@ def setup_jax_extension(
     else:
         cxx_flags.append("-g0")
 
+    if bool(int(os.getenv("NVTE_UB_WITH_MPI", "0"))):
+        assert os.getenv("MPI_HOME") is not None, (
+            "MPI_HOME=/path/to/mpi must be set when compiling with NVTE_UB_WITH_MPI=1!"
+        )
+        mpi_path = Path(os.getenv("MPI_HOME"))
+        include_dirs.append(mpi_path / "include")
+        cxx_flags.append("-DNVTE_UB_WITH_MPI")
+
     # Define TE/JAX as a Pybind11Extension
     from pybind11.setup_helpers import Pybind11Extension