MLP

MLP is the standard fully connected PhiJAX architecture. Its factory supports input normalization, coordinate features, dropout, independent output networks, and alternative dense-layer parameterizations.

Transformation order

The architecture applies optional operations in this order:

input normalization
-> periodic coordinate features
-> random Fourier features
-> hidden dense layers, activations, and dropout
-> output layer and optional output activation

Built-in architectures use the public ActivationName and InitializationName aliases. They also accept custom Activation and Initializer callables.

Option	Built-in names
activation	"relu", "leakyrelu", "gelu", "tanh", "sigmoid", "silu"
initialization	"kaiming_normal", "xavier_uniform", "trunc_normal" or None

one_mlp_per_output=True creates an independent scalar network per output. weight_norm and weight_factorization are mutually exclusive. Stochastic dropout requires deterministic=False and an explicit dropout_key.

Factory example

from functools import partial

from phijax.models import build_mlp

model_factory = partial(
    build_mlp,
    input_dim=2,
    output_dim=1,
    hidden=(128, 128, 128, 128),
    activation="tanh",
    input_norm=True,
)

See Layers for periodic features, random Fourier features, and factorized dense layers.

API reference

MLP [source]

MLP(
    input_dim: int,
    output_dim: int,
    hidden: Sequence[int] | None = (128,),
    activation: ActivationName | Activation | None = "silu",
    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,
    dropout: float = 0.0,
    initialization: InitializationName
    | Initializer
    | None = "xavier_uniform",
    initialization_kwargs: Mapping[str, Any] | None = None,
    one_mlp_per_output: bool = False,
    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 a dynamic multilayer perceptron with functional JAX state.

Attributes:

Name	Type	Description
periodic_embedding		Optional fixed periodic coordinate mapping.
embedding		Optional random Fourier feature mapping.
networks		One shared network or one scalar-output network per output.

Initialize a configurable MLP.

Parameters:

Name	Type	Description	Default
input_dim	int	Width of the input vector.	required
output_dim	int	Width of the output vector.	required
hidden	Sequence[int] | None	Width of each hidden layer. None or an empty sequence creates a direct input-to-output network.	(128,)
activation	ActivationName | Activation | None	Hidden activation name, callable, or None.	'silu'
activation_kwargs	Mapping[str, Any] | None	Keyword arguments bound to each 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
dropout	float	Hidden-activation dropout rate. Stochastic calls require an explicit key.	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
one_mlp_per_output	bool	Whether to construct an independent scalar network for each output.	False
weight_norm	bool	Whether to wrap dense layers with NNX weight normalization.	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 inputs with 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 while parameters remain float32.	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
TypeError	If compute_dtype cannot be resolved by JAX.
ValueError	If dimensions, names, dropout, normalization, or 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 of the form (..., slice(i, i + 1)) preserving the scalar output dimension.

Raises:

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

__call__ [source]

__call__(
    inputs: Array,
    input_mean: Array | None = None,
    input_std: Array | None = None,
    *,
    deterministic: bool = True,
    dropout_key: Array | None = None,
) -> jax.Array

Evaluate the 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
deterministic	bool	Whether to disable configured dropout.	True
dropout_key	Array | None	Explicit PRNG key required for stochastic dropout.	None

Returns:

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

Raises:

Type	Description
ValueError	If input width or normalization arrays are invalid, or stochastic dropout lacks a key.

build_mlp [source]

build_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 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 input 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:MLP.	{}

Returns:

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