Function make_mlp

test/helpers.py:54–55

(key)

@eqx.filter_jit
def make_mlp(key):
    return eqx.nn.MLP(2, 2, 2, 2, key=key)


@contextlib.contextmanager