waxmorph.jax.mlp.MLP

Contents

waxmorph.jax.mlp.MLP#

class waxmorph.jax.mlp.MLP(input_dim, output_dim, hidden_dim=128, num_layers=2, activation='silu', layer_norm=True, *, key)[source]#

Bases: Module

MLP with optional LayerNorm on its output.

num_layers counts linear maps; one layer is affine. Hidden layers support relu, silu, gelu, and tanh. Callers must supply a PRNG key; independent modules need independently split keys.

Raises:

  • TypeError – If num_layers is boolean or nonintegral.

  • ValueError – If num_layers < 1 or the activation is unsupported.

Parameters:

Examples

>>> mlp = MLP(3, 2, hidden_dim=4, num_layers=1, layer_norm=False, key=jax.random.PRNGKey(0))
>>> print(tuple(mlp(jnp.ones((5, 3))).shape))
(5, 2)
__call__(x)[source]#

Preserve arbitrary leading axes by vectorizing Equinox’s single-vector network.

Parameters:

x (Array)

Return type:

Array

Contents