JAX guide

mononet.jax uses Flax NNX — the object-oriented Flax API. Layers are flax.nnx.Module subclasses, fully compatible with jax.jit() and jax.grad(), and compose with flax.nnx.Sequential.

Install

pip install "mononet[jax]"

Public API

• mononet.jax.layers.MonoLinear — monotonic analogue of flax.nnx.Linear (non-decreasing in all inputs).

• mononet.jax.layers.MonoResidual — dual-gated monotone residual block, warm-started near identity.

• mononet.jax.layers.MonoInput — sign-flip layer encoding per-feature monotonicity directions.

mononet ships layers only — stack them yourself; there is no composed MonoMLP model.

Example

import jax.numpy as jnp
from flax import nnx
from mononet.jax import MonoInput, MonoLinear

# A monotonic MLP: non-decreasing in every input feature.
net = nnx.Sequential(
    MonoInput(1),                                    # +1 => up; -1 => down
    MonoLinear(4, 32, mode="switch", rngs=nnx.Rngs(0)),
    MonoLinear(32, 1, mode="switch", rngs=nnx.Rngs(1)),
)
y = net(jnp.zeros((8, 4)))                           # (8, 1), monotone in all inputs

The dense layers take an explicit rngs (flax.nnx.Rngs) for weight initialization. For per-feature monotonicity directions, pass a MonotonicityMask (a 1-D array of {-1, +1}) to MonoInput.

See also

• Concepts: monotonicity

• Layer reference

• Benchmarks