How to use the transonic.mpi function in transonic

codes_dependance,
            codes_dependance_classes,
            code_ext,
            special,
        ) = analysis_jit(self.get_source(), self.filename, backend_name)

        self.info_analysis = {
            "jitted_dicts": jitted_dicts,
            "codes_dependance": codes_dependance,
            "codes_dependance_classes": codes_dependance_classes,
            "special": special,
        }

        self.backend = backend = backends[backend_name]
        path_jit = mpi.Path(backend.jit.path_base)
        path_jit_class = mpi.Path(backend.jit.path_class)

        # TODO: check if these files have to be written here...
        # Write exterior code for functions
        for file_name, code in code_ext["function"].items():
            path_ext = path_jit / self.module_name.replace(".", os.path.sep)
            path_ext_file = path_ext / (file_name + ".py")
            write_if_has_to_write(path_ext_file, format_str(code), logger.info)

        # Write exterior code for classes
        for file_name, code in code_ext["class"].items():
            path_ext = path_jit_class / self.module_name.replace(".", os.path.sep)
            path_ext_file = path_ext / (file_name + ".py")
            write_if_has_to_write(path_ext_file, format_str(code), logger.info)

def block_until_avail(self, parallel=True):

        if mpi.rank == 0:
            if parallel:
                limit = self.limit_nb_processes
            else:
                limit = 1

            while len(self.processes) >= limit:
                time.sleep(self.deltat)
                self.processes = [
                    process
                    for process in self.processes
                    if process.is_alive_root()
                ]

        mpi.barrier(timeout=None)

module_name = get_module_name(frame)

    if backend_name is None:
        backend_name = get_backend_name_module(module_name)

    modules = modules_backends[backend_name]

    if module_name in modules:
        ts = modules[module_name]
        if (
            ts.is_transpiling != is_transpiling
            or ts._compile_at_import_at_creation != has_to_compile_at_import()
            or (
                hasattr(ts, "path_mod")
                and ts.path_backend.exists()
                and mpi.has_to_build(ts.path_backend, ts.path_mod)
            )
        ):
            ts = Transonic(frame=frame, reuse=False, backend=backend_name)
    else:
        ts = Transonic(frame=frame, reuse=False, backend=backend_name)

    return ts

def _load_old_header(self, path_backend_header):
        exports_old = None
        if mpi.rank == 0:
            with open(path_backend_header) as file:
                exports_old = [export.strip() for export in file.readlines()]
        exports_old = mpi.bcast(exports_old)
        return exports_old

def __init__(self, name, type_formatter):
        self.name = name
        self.name_capitalized = name.capitalize()
        self.type_formatter = type_formatter

        self.path_base = Path(path_root) / self.name / "__jit__"
        self.path_class = self.path_base.parent / "__jit_class__"

        if mpi.rank == 0:
            self.path_base.mkdir(parents=True, exist_ok=True)
            self.path_class.mkdir(parents=True, exist_ok=True)
        mpi.barrier()

if mpi.rank == 0:
            if parallel:
                limit = self.limit_nb_processes
            else:
                limit = 1

            while len(self.processes) >= limit:
                time.sleep(self.deltat)
                self.processes = [
                    process
                    for process in self.processes
                    if process.is_alive_root()
                ]

        mpi.barrier(timeout=None)

# for backend like numba
            if not self.compiling:
                backend_module = import_from_path(self.path_extension, name_mod)
                assert backend.check_if_compiled(backend_module)
                self.backend_func = getattr(backend_module, func_name)

        ext_files = None
        if mpi.rank == 0:
            glob_name_ext_file = (
                func_name + "_" + hex_src + "_*" + backend.suffix_extension
            )
            ext_files = list(
                mpi.PathSeq(path_backend).parent.glob(glob_name_ext_file)
            )
        ext_files = mpi.bcast(ext_files)

        if not ext_files:
            if has_to_compile_at_import() and _COMPILE_JIT:
                backenize_with_new_header()
            self.backend_func = None
        else:
            path_ext = max(ext_files, key=lambda p: p.stat().st_ctime)
            backend_module = import_from_path(path_ext, name_mod)
            self.backend_func = getattr(backend_module, func_name)

        # this is the function that will be called by the user
        @wraps(func)
        def type_collector(*args, **kwargs):

            if self.compiling:
                if not self.process.is_alive(raise_if_error=True):

path.name,
        ]

        words_command.extend(("-o", name_ext_file))

        words_command.extend(flags)

        cwd = path.parent

        self.block_until_avail(parallel)

        process = None
        if mpi.rank == 0:
            process = subprocess.Popen(words_command, cwd=cwd)

        process = mpi.ShellProcessMPI(process)

        if mpi.rank == 0:
            self.processes.append(process)
        return process