Utilities

Utilities support reproducibility, rank-aware logging, configuration display, task cleanup, shape normalization, and conversion between JAX and OmegaConf values.

Reproducibility

resolve_seed(None) selects a random unsigned 32-bit seed. The training entrypoint writes that concrete value back to the composed config and logs it before initialization. seed_everything seeds Python and NumPy and returns a JAX key folded with the current process index.

resolve_seed [source]

resolve_seed(seed: int | None) -> int

Resolve an optional root seed to a concrete unsigned 32-bit integer.

A configured integer is preserved for reproducibility. Passing None draws one seed from the operating system's cryptographic random source so it does not depend on Python or NumPy RNG state that this seed will subsequently initialize.

Parameters:

Name	Type	Description	Default
`seed`	`int \| None`	Configured root seed, or `None` to generate one randomly.	required

Returns:

Type	Description
`int`	Concrete unsigned 32-bit root seed.

Raises:

Type	Description
`TypeError`	If `seed` is neither an integer nor `None`, or is Boolean.
`ValueError`	If an integer seed lies outside the unsigned 32-bit range.

seed_everything [source]

seed_everything(
    seed: int, *, process_index: int | None = None
) -> jax.Array

Globally seed Python and NumPy and return an explicit process-local JAX key.

JAX has no implicit global random generator, so callers must still thread the returned key through model initialization, sampling, and functional training state. In distributed execution, host RNGs and the returned JAX key are deterministically separated by process index.

Parameters:

Name	Type	Description	Default
`seed`	`int`	Reproducible unsigned 32-bit root seed.	required
`process_index`	`int \| None`	Optional global JAX process index. Defaults to :func:`jax.process_index`.	`None`

Returns:

Type	Description
`Array`	Process-local typed JAX PRNG key derived from `seed`.

Raises:

Type	Description
`ValueError`	If `seed` or `process_index` is outside its valid range.

Python logging

RankedLogger [source]

RankedLogger(
    name: str = __name__,
    *,
    rank_zero_only: bool = False,
    extra: Mapping[str, object] | None = None,
)

Bases: LoggerAdapter

Standard-library logger adapter with optional rank-zero-only emission.

The adapter delegates to the logger registered under name, so handlers installed by Hydra remain responsible for formatting and routing records. When rank_zero_only is enabled, records created on nonzero JAX processes are discarded before they reach those handlers.

Initialize a JAX-aware logger adapter.

Parameters:

Name	Type	Description	Default
`name`	`str`	Name of the underlying standard-library logger.	`__name__`
`rank_zero_only`	`bool`	Whether to emit records only from JAX process zero.	`False`
`extra`	`Mapping[str, object] \| None`	Optional contextual values added to each emitted log record.	`None`

log [source]

log(
    level: int, msg: object, *args: object, **kwargs: Any
) -> None

Emit a record when its level and configured process policy allow it.

Parameters:

Name	Type	Description	Default
`level`	`int`	Standard-library logging level for the record.	required
`msg`	`object`	Message object passed to the underlying logger.	required
`*args`	`object`	Positional arguments forwarded to :meth:`logging.Logger.log`.	`()`
`**kwargs`	`Any`	Keyword arguments forwarded to :meth:`logging.Logger.log`.	`{}`

get_colorlogger [source]

get_colorlogger(
    name: str = "app", level: int = logging.INFO
) -> logging.Logger

Create or retrieve a colorized Hydra-style logger.

The logger owns one colorized stream handler and disables propagation so Hydra's root handler does not emit each message a second time. Repeated calls with the same name reuse the configured handler.

Parameters:

Name	Type	Description	Default
`name`	`str`	Logger name displayed in each record.	`'app'`
`level`	`int`	Minimum standard-library logging level to emit.	`INFO`

Returns:

Name	Type	Description
`Configured`	`Logger`	class:`logging.Logger` instance.

pad_keys [source]

pad_keys(
    mapping: Mapping[str, Any],
    prefix: str | None = None,
    postfix: str | None = None,
    exclude: str | Sequence[str] | None = None,
) -> dict[str, Any]

Add a prefix and postfix to selected mapping keys.

Parameters:

Name	Type	Description	Default
`mapping`	`Mapping[str, Any]`	Mapping whose string keys will be transformed.	required
`prefix`	`str \| None`	Text to prepend to keys, or `None` for no prefix.	`None`
`postfix`	`str \| None`	Text to append to keys, or `None` for no postfix.	`None`
`exclude`	`str \| Sequence[str] \| None`	One key or a sequence of keys to leave unchanged.	`None`

Returns:

Type	Description
`dict[str, Any]`	New dictionary containing the transformed keys and original values.

Configuration and task lifecycle

pre_hydra_routine must run before a decorated Hydra entrypoint so project-root discovery, environment defaults, and OmegaConf resolver registration happen before composition. task_wrapper logs failures and always runs registered finalizers, including optional W&B shutdown.

pre_hydra_routine [source]

pre_hydra_routine() -> None

Configure repository paths and environment variables before Hydra starts.

The repository root is discovered from this source file using pyproject.toml, exported as PROJECT_ROOT, and used to load an optional repository-level .env file. PhiJAX's process-global OmegaConf resolvers are registered before Hydra composes a config. The working directory and Python import path are unchanged.

extras [source]

extras(cfg: DictConfig) -> None

Apply optional warning and Rich-rendering utilities before a task starts.

Parameters:

Name	Type	Description	Default
`cfg`	`DictConfig`	Configuration containing an optional `extras` group and `paths.output_dir`.	required

print_config_tree [source]

print_config_tree(
    cfg: DictConfig,
    print_order: Sequence[str] = (
        "application",
        "data",
        "model",
        "callbacks",
        "logger",
        "trainer",
        "paths",
        "extras",
    ),
    resolve: bool = False,
    save_to_file: bool = False,
) -> None

Print a composed Hydra configuration as a Rich tree.

Parameters:

Name	Type	Description	Default
`cfg`	`DictConfig`	Configuration composed by Hydra.	required
`print_order`	`Sequence[str]`	Preferred order for top-level configuration fields. Present fields not listed here are appended in their configuration order.	`('application', 'data', 'model', 'callbacks', 'logger', 'trainer', 'paths', 'extras')`
`resolve`	`bool`	Whether to resolve interpolations before rendering configuration groups.	`False`
`save_to_file`	`bool`	Whether to save a plain-text rendering as `config_tree.log` in `cfg.paths.output_dir`.	`False`

Raises:

Type	Description
`ValueError`	If `save_to_file` is enabled but `cfg.paths.output_dir` is unavailable.

task_wrapper [source]

task_wrapper(
    task_func: Callable[[DictConfig], TaskReturn],
) -> Callable[[DictConfig], TaskReturn]

Decorate a Hydra task with assertion resolution and consistent lifecycle logging.

Lifecycle owners finalize stage-local resources normally. Longer-lived resources with idempotent cleanup, such as trainer checkpoint services and remote experiment loggers, may register a failure-safe finalizer through :func:register_task_finalizer; the wrapper invokes those finalizers in last-created-first-closed order.

Parameters:

Name	Type	Description	Default
`task_func`	`Callable[[DictConfig], TaskReturn]`	Task callable accepting the composed configuration as its only argument.	required

Returns:

Type	Description
`Callable[[DictConfig], TaskReturn]`	Wrapped callable that preserves the task's return value and re-raises failures with their original traceback.

register_task_finalizer [source]

register_task_finalizer(finalizer: TaskFinalizer) -> bool

Register an idempotent resource finalizer with the active task wrapper.

Parameters:

Name	Type	Description	Default
`finalizer`	`TaskFinalizer`	Callable accepting the terminal status `success`, `interrupted`, or `failed`.	required

Returns:

Type	Description
`bool`	Whether an active :func:`task_wrapper` accepted the finalizer.

Conversion and shapes

as_numpy [source]

as_numpy(value: Any) -> np.ndarray

Convert an array-like value to a host NumPy array.

JAX is imported only when it is installed and conversion is requested. This keeps configuration and evaluation utilities importable without a JAX runtime while ensuring device arrays are explicitly transferred to the host.

Parameters:

Name	Type	Description	Default
`value`	`Any`	Array-like value to convert. JAX arrays may reside on any supported device.	required

Returns:

Type	Description
`ndarray`	Host NumPy view or copy of `value`. JAX `bfloat16` arrays are promoted to `float32` because NumPy does not
`ndarray`	provide a native `bfloat16` dtype.

to_plain_container [source]

to_plain_container(
    value: Any, *, resolve: bool = True
) -> Any

Convert an OmegaConf container to ordinary Python containers.

Parameters:

Name	Type	Description	Default
`value`	`Any`	Value that may be an OmegaConf configuration container.	required
`resolve`	`bool`	Whether to resolve OmegaConf interpolations during conversion.	`True`

Returns:

Type	Description
`Any`	Plain Python containers for OmegaConf inputs, otherwise `value` unchanged.

as_tuple [source]

as_tuple(shape: int | Sequence[int]) -> tuple[int, ...]

Normalize a scalar or integer sequence shape to a tuple.

Parameters:

Name	Type	Description	Default
`shape`	`int \| Sequence[int]`	Scalar dimension or sequence of dimensions.	required

Returns:

Type	Description
`tuple[int, ...]`	Tuple containing all dimensions from `shape`.