Modified MLP

ModifiedMLP follows the gated architecture introduced for mitigating gradient-flow pathologies in PINNs. It creates two shallow coordinate encodings, U and V, then uses each activated hidden layer as an element-wise interpolation gate between them:

U = activation(dense_u(features))
V = activation(dense_v(features))
H = activation(dense(H))
H = H * U + (1 - H) * V

Where applicable, the architecture supports the same normalization, coordinate features, output names, precision, weight normalization, and random weight factorization as the standard MLP. It follows Wang, Teng, and Perdikaris (2021).

Factory example

from functools import partial

from phijax.models import build_modified_mlp

model_factory = partial(
    build_modified_mlp,
    input_dim=2,
    output_dim=1,
    hidden_dim=256,
    num_layers=4,
    activation="tanh",
    input_norm=True,
)

API reference

ModifiedMLP [source]

ModifiedMLP(
    input_dim: int,
    output_dim: int,
    hidden_dim: int = 256,
    num_layers: int = 4,
    activation: ActivationName | Activation = "tanh",
    activation_kwargs: Mapping[str, Any] | None = None,
    output_activation: ActivationName
    | Activation
    | None = None,
    output_activation_kwargs: Mapping[str, Any]
    | None = None,
    output_names: Sequence[str] | None = None,
    initialization: InitializationName
    | Initializer
    | None = "xavier_uniform",
    initialization_kwargs: Mapping[str, Any] | None = None,
    weight_norm: bool = False,
    weight_factorization: bool = False,
    weight_factorization_kwargs: Mapping[str, Any]
    | None = None,
    periodic_features: bool = False,
    periodic_features_kwargs: Mapping[str, Any]
    | None = None,
    fourier_features: bool = False,
    fourier_features_kwargs: Mapping[str, Any]
    | None = None,
    input_norm: bool = False,
    input_norm_eps: float = 1e-08,
    compute_dtype: Any = jnp.float32,
    parameter_dtype: Any = jnp.float32,
    output_dtype: Any = jnp.float32,
    *,
    rngs: Rngs,
)

Bases: Module

Implement the gated modified MLP architecture used in modern PINN training.

Two shallow latent encodings, U and V, are computed directly from the coordinate features. Every hidden layer then interpolates between these encodings using its activated output, improving gradient flow through deep PINNs.

Attributes:

Name	Type	Description
periodic_embedding		Optional fixed periodic coordinate mapping.
embedding		Optional random Fourier feature mapping.
gate_u		Dense layer producing the first shallow latent encoding.
gate_v		Dense layer producing the second shallow latent encoding.
layers		Gated hidden layers in evaluation order.
output_layer		Final linear prediction layer.

References

Wang, S., Teng, Y., and Perdikaris, P. (2021). Understanding and Mitigating Gradient Flow Pathologies in Physics-Informed Neural Networks. SIAM Journal on Scientific Computing, 43(5), A3055-A3081.

Initialize a modified MLP.

Parameters:

Name	Type	Description	Default
input_dim	int	Width of the physical input coordinate vector.	required
output_dim	int	Width of the predicted output vector.	required
hidden_dim	int	Shared width of latent gates and hidden layers.	256
num_layers	int	Number of gated hidden layers.	4
activation	ActivationName | Activation	Hidden activation name or callable.	'tanh'
activation_kwargs	Mapping[str, Any] | None	Keyword arguments bound to the hidden activation.	None
output_activation	ActivationName | Activation | None	Optional output activation name or callable.	None
output_activation_kwargs	Mapping[str, Any] | None	Keyword arguments bound to the output activation.	None
output_names	Sequence[str] | None	Optional unique names for scalar outputs.	None
initialization	InitializationName | Initializer | None	Dense-kernel initializer name, callable, or None for the NNX default.	'xavier_uniform'
initialization_kwargs	Mapping[str, Any] | None	Keyword arguments forwarded to the initializer factory.	None
weight_norm	bool	Whether to apply NNX weight normalization to dense layers.	False
weight_factorization	bool	Whether to use random weight-factorized dense layers.	False
weight_factorization_kwargs	Mapping[str, Any] | None	Options for :class:FactorizedDense, including mean, std, and use_bias.	None
periodic_features	bool	Whether to replace selected coordinates with fixed cosine and sine features.	False
periodic_features_kwargs	Mapping[str, Any] | None	Options for :class:PeriodicFeatures, including axes and frequencies.	None
fourier_features	bool	Whether to embed coordinates with trainable random Fourier features.	False
fourier_features_kwargs	Mapping[str, Any] | None	Options for :class:RandomFourierFeatures, including embed_dim and scale.	None
input_norm	bool	Whether to apply explicit per-input mean and standard deviation arrays.	False
input_norm_eps	float	Minimum standard deviation used during normalization.	1e-08
compute_dtype	Any	Data type used for model arithmetic.	float32
parameter_dtype	Any	Data type used to initialize trainable parameters.	float32
output_dtype	Any	Data type returned by the network.	float32
rngs	Rngs	NNX random streams used for parameter initialization.	required

Raises:

Type	Description
ValueError	If dimensions, names, normalization, activation, or dense-layer options are invalid.

slice_output [source]

slice_output(name: str) -> tuple[EllipsisType, slice]

Return an index tuple selecting one named scalar output.

Parameters:

Name	Type	Description	Default
name	str	Output name registered in output_names.	required

Returns:

Type	Description
tuple[EllipsisType, slice]	An index preserving the selected scalar output dimension.

Raises:

Type	Description
ValueError	If output names were not configured.
KeyError	If name is unknown.

coordinate_features [source]

coordinate_features(
    inputs: Array,
    input_mean: Array | None = None,
    input_std: Array | None = None,
) -> jax.Array

Apply normalization and configured coordinate embeddings.

Parameters:

Name	Type	Description	Default
inputs	Array	Coordinate array with final width input_dim.	required
input_mean	Array | None	Optional per-input normalization mean.	None
input_std	Array | None	Optional per-input normalization standard deviation.	None

Returns:

Type	Description
Array	Embedded coordinate features consumed by the gates and hidden layers.

Raises:

Type	Description
ValueError	If input width or normalization arrays are invalid.

__call__ [source]

__call__(
    inputs: Array,
    input_mean: Array | None = None,
    input_std: Array | None = None,
) -> jax.Array

Evaluate the gated modified MLP without mutating model state.

Parameters:

Name	Type	Description	Default
inputs	Array	Input array with final width input_dim.	required
input_mean	Array | None	Optional per-input normalization mean.	None
input_std	Array | None	Optional per-input normalization standard deviation.	None

Returns:

Type	Description
Array	Predictions with unchanged leading dimensions and final width output_dim.

build_modified_mlp [source]

build_modified_mlp(
    key: Array,
    input_dim: int,
    output_dim: int,
    input_mean: ArrayLike | None = None,
    input_std: ArrayLike | None = None,
    *,
    precision: PrecisionName
    | PrecisionPolicy
    | None = None,
    **model_kwargs: Any,
) -> InitializedModel

Build a normalized modified MLP through the generic initialized-model contract.

Parameters:

Name	Type	Description	Default
key	Array	JAX PRNG key for deterministic parameter initialization.	required
input_dim	int	Width of the physical input coordinate vector.	required
output_dim	int	Width of the predicted output vector.	required
input_mean	ArrayLike | None	Optional array-like per-coordinate input normalization mean.	None
input_std	ArrayLike | None	Optional array-like per-coordinate input normalization standard deviation.	None
precision	PrecisionName | PrecisionPolicy | None	Optional model precision policy.	None
**model_kwargs	Any	Architecture options forwarded to :class:ModifiedMLP.	{}

Returns:

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