waxmorph.jax.mlp

Equinox MLP used by the JAX graph-network simulator.

Classes

MLP(input_dim, output_dim[, hidden_dim, ...])

MLP with optional LayerNorm on its output.