Models

PhiJAX includes Flax NNX architectures, but equations and objectives only require a pure explicit-state callable. A ModelFactory delays initialization until the Trainer has prepared data and selected precision. It returns an InitializedModel containing an apply callable, explicit state, and optional summary.

Model contracts

ModelFactory [source]

Bases: Protocol

Build an explicit-state model after Trainer-owned data preparation.

Model factories receive normalization statistics and precision through keyword arguments so applications can bind architecture options independently with :class:functools.partial.

__call__ [source]

__call__(
    *,
    key: Array,
    input_mean: ArrayLike | None,
    input_std: ArrayLike | None,
    precision: PrecisionName | PrecisionPolicy,
) -> InitializedModel

Initialize a pure model application and its explicit state.

Parameters:

Name	Type	Description	Default
key	Array	Model-parameter initialization key.	required
input_mean	ArrayLike | None	Optional per-coordinate normalization mean supplied by the DataModule.	required
input_std	ArrayLike | None	Optional per-coordinate normalization standard deviation supplied by the DataModule.	required
precision	PrecisionName | PrecisionPolicy	Trainer precision policy.	required

Returns:

Type	Description
InitializedModel	Initialized pure model contract.

InitializedModel [source] dataclass

InitializedModel(
    apply: ModelApply,
    state: Any,
    summary: ModelSummaryFunction | None = None,
)

Bundle a pure model application with its explicit initialized state.

Attributes:

Name	Type	Description
apply	ModelApply	Pure callable accepting explicit model state and input arrays.
state	Any	Initialized differentiable model-state PyTree.
summary	ModelSummaryFunction | None	Optional callable rendering a summary for the explicit state.

Choose an architecture

Architecture	Use when
MLP	A conventional configurable coordinate network is sufficient
Modified MLP	Gated coordinate encodings may improve gradient flow
PirateNet	Training should gradually activate residual network depth
Layers	A custom architecture needs PhiJAX feature or dense-layer utilities

Each built-in architecture has a factory that returns InitializedModel. Bind architecture options with functools.partial and pass the resulting factory to PhiModule.

Custom NNX architectures

initialize_nnx_model() adapts an initialized Flax NNX module to InitializedModel. The application factory supplies the explicit key, normalization statistics, and precision policy expected by the Trainer. See Build a custom model for a complete factory example and a non-NNX checklist.

initialize_nnx_model [source]

initialize_nnx_model(
    model: Module,
    *,
    example_inputs: Array | None = None,
    call_kwargs: Mapping[str, Any] | None = None,
) -> InitializedModel

Adapt an initialized Flax NNX module to PhiJAX's explicit-state model contract.

The adapter separates the static NNX graph from its array state and captures only static call configuration in the returned application callable. Architecture factories remain responsible for constructing model with an explicit PRNG key and choosing representative example_inputs.

Parameters:

Name	Type	Description	Default
model	Module	Initialized NNX module to split into a graph definition and array state.	required
example_inputs	Array | None	Optional representative inputs used to provide a model-summary callable.	None
call_kwargs	Mapping[str, Any] | None	Optional keyword arguments bound to every model call, such as normalization statistics.	None

Returns:

Type	Description
InitializedModel	Pure model application, explicit NNX state, and an optional model-summary callable.

Model summaries

tabulate_nnx_model() creates the summary function used by built-in factories. Pass the result through InitializedModel.summary when adapting a custom NNX architecture.

tabulate_nnx_model [source]

tabulate_nnx_model(
    graphdef: Any,
    model_state: Any,
    *,
    example_inputs: Array,
    max_depth: int = -1,
    console_width: int = 120,
    compute_flops: bool = False,
    compute_vjp_flops: bool = False,
) -> str

Render a Flax NNX graph and explicit state with :func:flax.nnx.tabulate.

Parameters:

Name	Type	Description	Default
graphdef	Any	Static NNX graph definition paired with model_state.	required
model_state	Any	Explicit NNX variable state compatible with graphdef.	required
example_inputs	Array	Representative batched model inputs used to trace shapes.	required
max_depth	int	Maximum displayed module depth, or -1 for every level.	-1
console_width	int	Positive Rich console width in characters.	120
compute_flops	bool	Whether to estimate forward-pass floating-point operations.	False
compute_vjp_flops	bool	Whether to estimate reverse-pass floating-point operations.	False

Returns:

Type	Description
str	Rich-formatted model-summary table.

Raises:

Type	Description
ValueError	If max_depth is below -1 or console_width is not positive.