PhiModule

PhiModule connects a model factory to an objective. It describes the computation that the Trainer initializes and runs. Optimizer, loss-balancer, and PRNG state remain outside the module in TrainState.

Use PhiModule for the standard objective-backed workflow. Subclass BasePhiModule only when an application needs a different numerical contract.

Basic use

Create an uninitialized module blueprint, fit it, and use the bound module returned in FitResult for prediction:

module = PhiModule(model_factory, objective, name="Heat PINN")

result = trainer.fit(
    module,
    datamodule=data_module,
    optimizer=optimizer,
    seed=0,
)
predictions = trainer.predict(result, datamodule=data_module)

Trainer.fit() supplies the model key, precision policy, and optional input statistics. It does not modify the original blueprint. See the heat-equation quickstart for a complete runnable example.

Responsibilities

Component	Responsibility
`PhiModule`	Model application, objective losses, prediction, and metric routing
Model factory	Model initialization and explicit model state
Objective	Named unweighted losses and required DataModule batch keys
`Trainer`	Runtime assembly, hooks, device placement, logging, and checkpointing
`TrainState`	Model, optimizer, balancer, step, precision, and explicit PRNG arrays
Loss balancer	Combining named losses and scheduling adaptive weight updates

This separation keeps the compiled computation functional. A module never owns mutable optimizer or balancer state.

Module blueprints

The standard PhiModule accepts either a lazy ModelFactory or an initialized InitializedModel, together with an Objective.

The objective supplies ordered loss_names and batch_keys. The Trainer uses them to create default static loss weights and request matching batches from the DataModule. During fitting, prepare_model() creates a shallow bound copy containing the pure model application and model summary function. The resulting bound module is available as result.module.

Logging metrics

The compiled training_step() returns numerical losses and diagnostics. The standard PhiModule.on_train_batch_end() uses self.log() to route every scalar to configured experiment loggers. By default, these values also appear in the progress bar:

train/loss;
each train/loss/<name> objective loss;
each train/weight/<name> balancer weight.

Override a destination after calling the standard hook:

def on_train_batch_end(self, model_state, context):
    model_state, metrics = super().on_train_batch_end(model_state, context)
    self.log("train/loss/pde/heat", metrics["train/loss/pde/heat"], prog_bar=False)
    self.log("train/weight/pde/heat", metrics["train/weight/pde/heat"], prog_bar=False)
    self.log("train/residual/mean_abs", metrics["train/residual/mean_abs"], prog_bar=True)
    return model_state, metrics

self.log() records the device value without writing immediately. The Trainer converts it only at the configured logging or progress-refresh cadence. The method is available during on_train_batch_end() and must not be called from inside a JAX transformation. See Logging and monitoring for logger and progress-bar behavior.

Custom modules

Subclass BasePhiModule when the standard model-factory and objective composition is not sufficient. Implement loss_names, forward(), and training_step(). Implement residual_stream() when an adaptive loss balancer needs raw residuals or output streams.

class HeatModule(BasePhiModule):
    @property
    def loss_names(self) -> tuple[str, ...]:
        return ("initial/u", "boundary/u", "pde/heat")

    def forward(self, model_state, inputs):
        return self.model_apply(model_state, inputs)

    def training_step(self, model_state, batches):
        return self.objective.losses(self.model_apply, model_state, batches)

A fully bound custom module uses Trainer.fit_state() with an explicit TrainingPlan and TrainState. Keep optimizer and loss-balancer selection outside the module.

Lifecycle and hooks

Callbacks run in declaration order before the matching module hook. The main fit ordering is:

callbacks.setup                 -> module.setup
callbacks.on_fit_start          -> module.on_fit_start
callbacks.on_train_batch_start  -> module.on_train_batch_start
compiled training step
callbacks.on_train_batch_end    -> module.on_train_batch_end
callbacks.on_fit_end            -> module.on_fit_end
callbacks.teardown              -> module.teardown

Module hooks run on the Python host. Training hooks return explicit replacement values and must preserve the PyTree structure and fixed shapes expected by compiled functions. Prediction hooks observe state; override predict_step() to transform output batches before collection or artifact writing.

See Trainer and module hooks for the complete fit, prediction, interruption, and exception lifecycles.

API reference

Standard module

PhiModule [source]

PhiModule(
    model: ModelFactory | InitializedModel,
    objective: Objective,
    *,
    name: str = "PINN",
)

Bases: BasePhiModule

Describe a model factory and objective before Trainer-owned initialization.

The user-facing instance is an immutable-in-practice blueprint. :class:phijax.Trainer calls :meth:prepare_model after the DataModule has exposed normalization statistics, producing a shallow bound copy with a pure model application and explicit model state. The original blueprint is never mutated.

Initialize an unbound module blueprint.

Parameters:

Name	Type	Description	Default
`model`	`ModelFactory \| InitializedModel`	Lazy model factory or an existing :class:`phijax.InitializedModel`.	required
`objective`	`Objective`	Objective producing named unweighted scalar losses.	required
`name`	`str`	Non-empty human-readable module name.	`'PINN'`

Raises:

Type	Description
`TypeError`	If `model` cannot initialize an explicit-state model.
`ValueError`	If `name` is empty or the objective exposes invalid loss names.

loss_names `property`

loss_names: tuple[str, ...]

Return stable objective loss names.

Returns:

Type	Description
`tuple[str, ...]`	Ordered unique objective loss names.

batch_keys `property`

batch_keys: tuple[str, ...]

Return objective batch names in stable declaration order.

Returns:

Type	Description
`tuple[str, ...]`	Ordered unique batch names inferred from the objective.

prepare_model [source]

prepare_model(
    *,
    key: Array,
    input_mean: ArrayLike | None,
    input_std: ArrayLike | None,
    precision: PrecisionPolicy,
) -> tuple[PhiModule, Any]

Initialize and bind the model without mutating this blueprint.

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`	`PrecisionPolicy`	Trainer precision policy.	required

Returns:

Type	Description
`tuple[PhiModule, Any]`	Bound shallow module copy and its explicit initialized model state.

Raises:

Type	Description
`TypeError`	If the configured factory does not return :class:`phijax.InitializedModel`.

summarize_model [source]

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

Render the configured model summary.

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 \| None`	Rendered model summary, or `None` when no summary provider was supplied.

on_train_batch_end [source]

on_train_batch_end(
    model_state: Any, context: PhiModuleContext
) -> tuple[Any, Mapping[str, Any]]

Declare the default logger and progress destinations for completed metrics.

Every scalar is sent to experiment loggers. The total train/loss, individual train/loss/<name> losses, and train/weight/<name> balancer weights are displayed in that order by default. Array diagnostics remain available to callbacks without an implicit reduction or host transfer.

Parameters:

Name	Type	Description	Default
`model_state`	`Any`	Updated explicit model state.	required
`context`	`PhiModuleContext`	Post-update context containing compiled metrics and diagnostics.	required

Returns:

Type	Description
`tuple[Any, Mapping[str, Any]]`	Unchanged model state and complete metric mapping.

Custom module contract

BasePhiModule [source]

BasePhiModule(*, name: str = 'PINN')

Bases: _ModuleHooks, ABC

Define the application-facing computation and lifecycle contract used by :class:Trainer.

A module owns model application and objective behavior but never owns optimizer or loss-balancer logic. Mutable arrays remain in :class:phijax.training.TrainState, while training hooks return explicit model-state and batch replacements to the trainer. Prediction lifecycle hooks are observational; override :meth:predict_step to change outputs. Hooks execute on the Python host and must preserve JAX PyTree structures expected by compiled steps. The trainer invokes each callback hook before the matching module hook to follow Lightning's lifecycle convention.

Initialize the base module identity.

Parameters:

Name	Type	Description	Default
`name`	`str`	Non-empty human-readable module name.	`'PINN'`

Raises:

Type	Description
`ValueError`	If `name` is empty or whitespace-only.

loss_names `abstractmethod` `property`

loss_names: Sequence[str]

Return stable names for the unweighted losses produced by :meth:training_step.

Returns:

Type	Description
`Sequence[str]`	Ordered unique loss names used to initialize an independently selected balancer.

batch_keys `property`

batch_keys: Sequence[str]

Return batch names required by the module's training computation.

Advanced modules used with :meth:phijax.Trainer.fit_state may leave this empty when their explicit :class:phijax.TrainingPlan already defines the batch routing.

Returns:

Type	Description
`Sequence[str]`	Ordered unique DataModule batch names, or an empty sequence when they cannot be inferred.

forward [source] `abstractmethod`

forward(model_state: Any, inputs: Array) -> jax.Array

Define the explicit-state forward computation.

Parameters:

Name	Type	Description	Default
`model_state`	`Any`	Explicit model parameter PyTree.	required
`inputs`	`Array`	One model input point.	required

Returns:

Type	Description
`Array`	Model output for `inputs`.

training_step [source] `abstractmethod`

training_step(
    model_state: Any, batches: NamedBatches
) -> Mapping[str, jax.Array] | TrainingOutput

Evaluate named unweighted losses for one compiled training step.

Parameters:

Name	Type	Description	Default
`model_state`	`Any`	Explicit differentiable model parameter PyTree.	required
`batches`	`NamedBatches`	Fixed-structure named objective batches.	required

Returns:

Type	Description
`Mapping[str, Array] \| TrainingOutput`	Scalar losses keyed by :attr:`loss_names`, optionally accompanied by compiled diagnostics.

log [source]

log(
    name: str,
    value: Any,
    *,
    logger: bool = True,
    prog_bar: bool = False,
) -> None

Route one host-side batch metric to loggers or progress displays.

This method is available only while :class:phijax.Trainer executes :meth:on_train_batch_end. It records device values without transferring them to the host; the Trainer performs conversion at its configured logging or display cadence.

Parameters:

Name	Type	Description	Default
`name`	`str`	Non-empty complete metric name.	required
`value`	`Any`	Scalar host or device value. Arrays are accepted only when both destinations are disabled.	required
`logger`	`bool`	Whether configured experiment loggers receive the metric.	`True`
`prog_bar`	`bool`	Whether progress callbacks display the metric.	`False`

Raises:

Type	Description
`RuntimeError`	If called outside :meth:`on_train_batch_end`.
`TypeError`	If `name`, `logger`, or `prog_bar` has an invalid type.
`ValueError`	If `name` is empty or a routed value is not scalar.

format_training_metrics [source]

format_training_metrics(
    total_loss: Array,
    losses: Mapping[str, Array],
    diagnostics: Mapping[str, Array],
) -> dict[str, jax.Array]

Format compiled training values for progress bars and experiment loggers.

The combined loss becomes train/loss, each unweighted objective value becomes train/loss/<name>, and diagnostics become train/<name>. The host Trainer logs every scalar by default and displays only train/loss unless :meth:log overrides those destinations.

Parameters:

Name	Type	Description	Default
`total_loss`	`Array`	Scalar objective after external loss balancing.	required
`losses`	`Mapping[str, Array]`	Named unweighted scalar losses returned by :meth:`training_step`.	required
`diagnostics`	`Mapping[str, Array]`	Named scalar balancer, precision, or optimizer diagnostics.	required

Returns:

Type	Description
`dict[str, Array]`	Stable flat metric mapping suitable for JIT output and host logging.

residual_stream [source]

residual_stream(
    name: str, model_state: Any, batches: NamedBatches
) -> jax.Array

Evaluate one residual stream for an adaptive external balancer.

Parameters:

Name	Type	Description	Default
`name`	`str`	One entry from :attr:`loss_names`.	required
`model_state`	`Any`	Explicit differentiable model parameter PyTree.	required
`batches`	`NamedBatches`	Fixed-structure named residual batches.	required

Returns:

Type	Description
`Array`	Raw residual or model-output stream.

Raises:

Type	Description
`NotImplementedError`	If the module does not expose residual streams.

summarize_model [source]

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

Render the underlying network architecture when the module supports summaries.

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 \| None`	Rendered model summary, or `None` when no summary provider is configured.

predict_step [source]

predict_step(
    model_state: Any, batch: ArrayMapping
) -> jax.Array

Evaluate one prediction batch.

Parameters:

Name	Type	Description	Default
`model_state`	`Any`	Explicit model parameter PyTree.	required
`batch`	`ArrayMapping`	Prediction batch containing model coordinates under `inputs`.	required

Returns:

Type	Description
`Array`	Batched model outputs.

Hook context

PhiModuleContext [source] `dataclass`

PhiModuleContext(step: int, metrics: Mapping[str, Any])

Expose module-specific host-loop progress without optimizer or balancer state.

Attributes:

Name	Type	Description
`step`	`int`	Current completed optimizer step as a Python integer.
`metrics`	`Mapping[str, Any]`	Most recent device or host metric mapping.

PhiModule

Basic use

Responsibilities

Module blueprints

Logging metrics

Custom modules

Lifecycle and hooks

API reference

Standard module

PhiModule [source]

loss_names property

batch_keys property

prepare_model [source]

summarize_model [source]

on_train_batch_end [source]

Custom module contract

BasePhiModule [source]

loss_names abstractmethod property

batch_keys property

forward [source] abstractmethod

training_step [source] abstractmethod

log [source]

format_training_metrics [source]

residual_stream [source]

summarize_model [source]

predict_step [source]

Hook context

PhiModuleContext [source] dataclass

loss_names `property`

batch_keys `property`

loss_names `abstractmethod` `property`

batch_keys `property`

forward [source] `abstractmethod`

training_step [source] `abstractmethod`

PhiModuleContext [source] `dataclass`