def METHOD_NAME(local_jit, *args, **kwargs):
    """A decorator for jit compilation.
    Use this function as a decorator to any jit compiler function. In a
    parallel run, this function will first call the jit compilation
    function on the first process. When this is done, and the module is
    in the cache, it will call the jit compiler on the remaining
    processes, which will then use the cached module.
    """
    @functools.wraps(local_jit)
    def mpi_jit(comm, *args, **kwargs):
        # Just call JIT compiler when running in serial
        if comm.size == 1:
            return local_jit(*args, **kwargs)
        # Default status (0 == ok, 1 == fail)
        status = 0
        # Compile first on process 0
        root = comm.rank == 0
        if root:
            try:
                output = local_jit(*args, **kwargs)
            except Exception as e:
                status = 1
                error_msg = str(e)
        # TODO: This would have lower overhead if using the dijitso.jit
        # features to inject a waiting callback instead of waiting out
        # here. That approach allows all processes to first look in the
        # cache, introducing a barrier only on cache miss. There's also
        # a sketch in dijitso of how to make only one process per
        # physical cache directory do the compilation.
        # Wait for the compiling process to finish and get status
        # TODO: Would be better to broadcast the status from root but
        # this works.
        global_status = comm.allreduce(status, op=MPI.MAX)
        if global_status == 0:
            # Success, call jit on all other processes (this should just
            # read the cache)
            if not root:
                output = local_jit(*args, **kwargs)
        else:
            # Fail simultaneously on all processes, to allow catching
            # the error without deadlock
            if not root:
                error_msg = "Compilation failed on root node."
            raise RuntimeError(f"Failed just-in-time compilation of form: {error_msg}")
        return output
    # Return the decorated jit function
    return mpi_jit