API reference

The API reference lists PhiJAX's supported Python imports. Signatures, parameters, return values, and exceptions come from source docstrings. Each page also explains ownership, array shapes, and lifecycle behavior. Optional configuration integration is documented separately.

Stability

Names re-exported by phijax or a package __init__.py are the public API. Private names beginning with _ are implementation details. Domain applications and their configuration trees live in separate project repositories.

PhiJAX is pre-1.0, so configuration names and Python interfaces may still change between minor releases. Migration notes should accompany changes to public imports or required config paths.

The installed release is available as phijax.__version__.

Core conventions

• Bound model, optimizer, balancer, PRNG, and step state are explicit JAX PyTrees.
• PhiModule blueprints retain model factories and objectives without mutable initialized parameters.
• A model application callable maps (model_state, point) to one output vector.
• Host data uses NumPy until the trainer-selected strategy performs device placement.
• Training arrays have fixed shapes and stable PyTree structures inside compiled computations.
• Objective, balancer, metric, callback, and checkpoint alignment uses ordered loss names.
• Every random operation accepts or derives an explicit JAX key.

Common types

Type	Contract
ArrayMapping	Field names mapped to JAX arrays
NamedBatches	Objective batch keys mapped to array mappings
ModelApply	Explicit-state single-point model callable
ResidualGroup	Residual arrays reduced into one scalar loss
ResidualGroups	Residual groups aligned with static equation names
ResidualStream	Literal "residual" or "output"
ModelSummaryFunction	Callable rendering an explicit model state
ResidualFunction	Configured equation callable consumed by an objective term
JaxDevice	Structural protocol for the stable device metadata used by strategies

Package-specific public aliases include:

• BatchSize, DevicePool, ArrayFormat, and DataStage in phijax.data;
• CallbackContext, MonitorMode, and LoggingInterval in phijax.callbacks;
• Activation, ActivationName, Initializer, and InitializationName in phijax.models; and
• Accelerator, DeviceSelection, PrecisionMode, PrecisionAlias, PrecisionName, and MatmulPrecision in phijax.training.

JaxDevice [source]

Bases: Protocol

Describe the stable JAX device metadata used by PhiJAX orchestration.

JAX exposes its concrete device class dynamically, which some static type checkers reject in annotations. This protocol keeps PhiJAX independent of private :mod:jaxlib implementation modules while accepting real JAX devices.

Attributes:

Name	Type	Description
platform	str	Generic device platform such as cpu, gpu, or tpu.
process_index	int	JAX process that owns the device.
id	int	Backend-local device identifier.

ModelSummaryFunction [source]

Bases: Protocol

Define a model-state summary callable independent of a concrete neural-network library.

__call__ [source]

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

Render a model summary for explicit model state.

Parameters:

Name	Type	Description	Default
model_state	Any	Explicit model parameter and variable state.	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	Rendered model-summary table.

ResidualFunction [source]

Bases: Protocol

Define one configured equation callable that produces grouped residual streams.

__call__ [source]

__call__(
    model_apply: ModelApply,
    model_state: Any,
    batch: ArrayMapping,
    *,
    stream: ResidualStream = "residual",
) -> ResidualGroups

Evaluate residual groups or an explicitly supported alternative stream.

Parameters:

Name	Type	Description	Default
model_apply	ModelApply	Pure explicit-state model application callable.	required
model_state	Any	Differentiable model parameter PyTree.	required
batch	ArrayMapping	Fixed-structure arrays required by the configured equation.	required
stream	ResidualStream	Stream representation requested by the objective.	'residual'

Returns:

Type	Description
ResidualGroups	Residual groups aligned with the owning objective term's loss names.

Import map

from phijax import DataModule, ModelFactory, PhiModule, Trainer
from phijax.balancers import BalancerUpdatePlan, ExactNTKBalancer, GradNormBalancer, StaticLossBalancer
from phijax.callbacks import EarlyStopping, ModelCheckpoint, ProgressBar, RichModelSummary, RichProgressBar
from phijax.data import HostPool, NamedBatchSource, PhiDataModule, create_sampler
from phijax.evaluation import EvaluationResult, RegressionEvaluator
from phijax.equations import burgers_1d, polar_navier_stokes, residual_equation
from phijax.models import MLP, ModifiedMLP, PirateNet, build_mlp, build_modified_mlp, build_pirate_net
from phijax.models import initialize_nnx_model
from phijax.objectives import CompositeObjective, ResidualTerm
from phijax.training import OrbaxCheckpointIO, TrainState, TrainingPlan
from phijax.utils import RankedLogger, resolve_seed, seed_everything

Use the navigation to open a focused reference page. Start with the Trainer, PhiModule, Data, and Models pages for the main runtime contracts.