def METHOD_NAME(parity, symmetries, lattice, mode):
    g, hi, perms = _setup_symm(symmetries, N=3, lattice=lattice)
    ma = nk.models.GCNN(
        symmetries=perms,
        mode=mode,
        shape=tuple(g.extent),
        layers=2,
        features=2,
        parity=parity,
        bias_init=uniform(),
    )
    pars = ma.init(nk.jax.PRNGKey(), hi.random_state(nk.jax.PRNGKey(), 1))
    v = hi.random_state(jax.random.PRNGKey(0), 3)
    # vals = [ma.apply(pars, v[..., p]) for p in np.asarray(perms)]
    # code below implements the commented line above, but is vectorised
    v = v[..., np.asarray(perms)].transpose(1, 0, 2)
    v = v.reshape(len(perms) * 3, g.n_nodes)
    vals = ma.apply(pars, v).reshape(len(perms), 3)
    for val in vals:
        assert jnp.allclose(val, vals[0])