# make sure to import mpi4py before anything else
# this calls MPI_Init and registers mpi4py's atexit handler
from mpi4py import MPI

# check version of JAX and jaxlib
from .jax_compat import check_jax_version

check_jax_version()

# this registers our custom XLA functions
from . import xla_bridge  # noqa: E402

# register atexit handler to flush JAX buffers
import atexit  # noqa: E402


def flush():
    """Wait for all pending XLA operations"""
    import jax

    jax.effects_barrier()


atexit.register(flush)

# import public API
from .collective_ops.allgather import allgather  # noqa: E402
from .collective_ops.allreduce import allreduce  # noqa: E402
from .collective_ops.alltoall import alltoall  # noqa: E402
from .collective_ops.barrier import barrier  # noqa: E402
from .collective_ops.bcast import bcast  # noqa: E402
from .collective_ops.gather import gather  # noqa: E402
from .collective_ops.recv import recv  # noqa: E402
from .collective_ops.reduce import reduce  # noqa: E402
from .collective_ops.scan import scan  # noqa: E402
from .collective_ops.scatter import scatter  # noqa: E402
from .collective_ops.send import send  # noqa: E402
from .collective_ops.sendrecv import sendrecv  # noqa: E402

from .utils import has_cuda_support, has_sycl_support  # noqa: E402

# sanitize namespace
del check_jax_version, xla_bridge, MPI, atexit, flush

__all__ = [
    "allgather",
    "allreduce",
    "alltoall",
    "barrier",
    "bcast",
    "gather",
    "recv",
    "reduce",
    "scan",
    "scatter",
    "send",
    "sendrecv",
    "has_cuda_support",
    "has_sycl_support",
]
