I understand internal functions are prefixed with '_' to indicate they are helper/internal functions. It also helps with tooling etc. But I find some functions with just '_' as their name. Can't even find where they are called from. e.g., from https://github.com/jax-ml/jax/blob/...2decc86/jax/_src/pallas/mosaic/helpers.py#L72 def run_on_first_core(core_axis_name: str): """Runs a function on the first core in a given axis.""" num_cores = jax.lax.axis_size(core_axis_name) if num_cores == 1: return lambda f: f() def wrapped(f): core_id = jax.lax.axis_index(core_axis_name) @pl_helpers.when(core_id == 0) @functools.wraps(f) def _(): ## How is this called? return f() return wrapped There are several of them in an internal code base but here are some references https://github.com/search?q=repo:jax-ml/jax def _():&type=code https://github.com/jax-ml/jax/blob/...ecc86/docs/pallas/tpu/distributed.ipynb#L1125 Continue reading...
