whestbench.
API

sample_mlp

Sample a random MLP with He-initialized weight matrices.

function · source

sample_mlp(width: 'int', depth: 'int', rng: 'Optional[fnp.random.Generator]' = None, *, seed: 'int' = 0) -> 'MLP'
Sample a random MLP with He-initialized weight matrices.

Each weight matrix has shape ``(width, width)`` with entries drawn from
``N(0, 2/width)`` (He initialization for ReLU networks).

Args:
    width: Neuron count per layer.
    depth: Number of weight matrices.
    rng: Optional flopscope RNG for weight sampling. If None, a fresh
        unseeded generator is used.
    seed: Per-MLP grader-supplied seed to attach to the returned MLP for
        estimator consumption. This does NOT control weight sampling
        (which uses ``rng``); it's a separate stream the estimator may
        consume for its own randomness.