import inspect
import logging

import jax


def log_for_0(msg, *args, level=logging.INFO):
    """Log only on the first process (process_index == 0)."""
    if jax.process_index() != 0:
        return
    caller_module = inspect.currentframe().f_back.f_globals.get("__name__", __name__)
    logging.getLogger(caller_module).log(level, msg, *args)