PirateNet

PirateNet implements the Physics-Informed Residual Adaptive Network from Wang et al. (2024). Each block contains three dense layers, two modified-MLP gates, and a trainable scalar alpha:

nonlinear = gated_three_layer_block(features, U, V)
features = alpha * nonlinear + (1 - alpha) * features

The default alpha=0 makes every block an identity at initialization. Training increases the effective depth by learning each coefficient. Random Fourier features are enabled by default and must produce hidden_dim features. Because PhiJAX pairs cosine and sine features, the default samples embed_dim=hidden_dim/2 frequencies. Random weight factorization can be enabled separately.

The factory uses standard Glorot initialization. It does not apply the paper's optional physics-informed least-squares initialization because that method depends on training data.

Factory example

from functools import partial

from phijax.models import build_pirate_net

model_factory = partial(
    build_pirate_net,
    input_dim=2,
    output_dim=1,
    hidden_dim=256,
    num_blocks=4,
    activation="tanh",
    nonlinearity=0.0,
    input_norm=True,
    fourier_features_kwargs={"scale": 1.0},
    weight_factorization=True,
    weight_factorization_kwargs={"mean": 0.5, "std": 0.1},
)

API reference

PirateBlock [source]

PirateBlock(
    hidden_dim: int,
    activation: Activation,
    *,
    nonlinearity: float = 0.0,
    kernel_init: Initializer | None = None,
    weight_factorization: bool = False,
    weight_factorization_kwargs: Mapping[str, Any]
    | None = None,
    compute_dtype: Any = jnp.float32,
    parameter_dtype: Any = jnp.float32,
    rngs: Rngs,
)

Bases: Module

Apply one gated PirateNet block with an adaptive residual connection.

Attributes:

Name	Type	Description
`layers`		Three hidden-width dense layers.
`alpha`		Trainable scalar controlling the block's nonlinear contribution.

Initialize an adaptive residual block.

Parameters:

Name	Type	Description	Default
`hidden_dim`	`int`	Shared width of block inputs, gates, and hidden activations.	required
`activation`	`Activation`	Point-wise activation callable.	required
`nonlinearity`	`float`	Initial value of the trainable residual coefficient `alpha`.	`0.0`
`kernel_init`	`Initializer \| None`	Optional dense-kernel initializer.	`None`
`weight_factorization`	`bool`	Whether to use random weight-factorized dense layers.	`False`
`weight_factorization_kwargs`	`Mapping[str, Any] \| None`	Options forwarded to :class:`FactorizedDense`.	`None`
`compute_dtype`	`Any`	Data type used for block arithmetic.	`float32`
`parameter_dtype`	`Any`	Data type used to initialize trainable parameters.	`float32`
`rngs`	`Rngs`	NNX random streams used for parameter initialization.	required

Raises:

Type	Description
`ValueError`	If `hidden_dim` is not positive or `nonlinearity` is not finite.

call [source]

__call__(
    inputs: Array, gate_u: Array, gate_v: Array
) -> jax.Array

Evaluate the gated dense stack and adaptive identity interpolation.

Parameters:

Name	Type	Description	Default
`inputs`	`Array`	Block inputs with final width `hidden_dim`.	required
`gate_u`	`Array`	First shallow latent gate with the same shape as `inputs`.	required
`gate_v`	`Array`	Second shallow latent gate with the same shape as `inputs`.	required

Returns:

Type	Description
`Array`	Interpolation between the block inputs and nonlinear block output.

PirateNet [source]

PirateNet(
    input_dim: int,
    output_dim: int,
    hidden_dim: int = 256,
    num_blocks: int = 2,
    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,
    nonlinearity: float | Sequence[float] = 0.0,
    initialization: InitializationName
    | Initializer
    | None = "xavier_uniform",
    initialization_kwargs: Mapping[str, Any] | None = None,
    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 = True,
    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 a Physics-Informed Residual Adaptive Network.

PirateNet embeds coordinates into hidden_dim features, constructs two shallow latent gates, and applies a stack of three-layer adaptive residual blocks. Each block's scalar coefficient is normally initialized to zero, making every block an exact identity map at initialization before progressively learning a nonlinear contribution.

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 gate.
`gate_v`		Dense layer producing the second shallow latent gate.
`blocks`		Adaptive residual blocks in evaluation order.
`output_layer`		Final linear prediction layer.

References

Wang, S., Li, B., Chen, Y., and Perdikaris, P. (2024). PirateNets: Physics-informed Deep Learning with Residual Adaptive Networks. Journal of Machine Learning Research, 25(402), 1-51.

Initialize a PirateNet architecture.

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 coordinate embeddings, gates, and residual blocks.	`256`
`num_blocks`	`int`	Number of adaptive residual blocks, each containing three dense layers.	`2`
`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`
`nonlinearity`	`float \| Sequence[float]`	Initial adaptive coefficient shared by all blocks or one value per block.	`0.0`
`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_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 use a trainable random Fourier coordinate embedding.	`True`
`fourier_features_kwargs`	`Mapping[str, Any] \| None`	Options for :class:`RandomFourierFeatures`. `embed_dim` counts sampled frequencies, so the resulting feature width is `2 * embed_dim` and must equal `hidden_dim`.	`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 while parameters remain separately configurable.	`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, embeddings, output names, nonlinearities, or normalization 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

Transform physical coordinates into fixed-width residual features.

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`	Coordinate features with final width `hidden_dim`.

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 PirateNet 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_pirate_net [source]

build_pirate_net(
    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 normalized PirateNet 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:`PirateNet`.	`{}`

Returns:

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

PirateNet

Factory example

API reference

PirateBlock [source]

__call__ [source]

PirateNet [source]

slice_output [source]

coordinate_features [source]

__call__ [source]

build_pirate_net [source]

call [source]

call [source]