def METHOD_NAME():
    """Tests the Jax Blackjack env with vmap."""
    env = BlackjackFunctional()
    num_envs = 10
    rng, *split_rng = jrng.split(
        jrng.PRNGKey(0), num_envs + 1
    )  # this plus 1 is important because we want
    # num_envs subkeys and a main entropy source key which necessitates an additional key
    env.transform(jax.vmap)
    env.transform(jax.jit)
    state = env.initial(jnp.array(split_rng))
    env.action_space.seed(0)
    for t in range(10):
        obs = env.observation(state)
        action = jnp.array([env.action_space.sample() for _ in range(num_envs)])
        # if isinstance(env.action_space, Discrete):
        #     action = action.reshape((num_envs, 1))
        rng, *split_rng = jrng.split(rng, num_envs + 1)
        next_state = env.transition(state, action, jnp.array(split_rng))
        terminal = env.terminal(next_state)
        reward = env.reward(state, action, next_state)
        assert len(next_state) == len(state)
        # assert next_state.dtype == jnp.float32
        assert reward.shape == (num_envs,)
        assert reward.dtype == jnp.float32
        assert terminal.shape == (num_envs,)
        assert terminal.dtype == bool
        assert isinstance(obs, jax.Array)
        assert obs[0].dtype == jnp.int32
        assert obs[1].dtype == jnp.int32
        assert obs[2].dtype == jnp.int32
        state = next_state