def METHOD_NAME(n, device_array):
    if device_array:
        n = jnp.asarray(n)
    p = np.array([0.5, 0.5])
    x = multinomial(random.PRNGKey(0), p, n)
    assert x.shape == jnp.shape(n) + jnp.shape(p)
    assert_allclose(x.sum(-1), n)