def METHOD_NAME(
    jax_dist_getter, nb_mixtures, mixing_batch_shape, component_batch_shape, same_family
):
    # Create mixture
    mixing_probabilities = jnp.ones(nb_mixtures) / nb_mixtures
    for i, s in enumerate(mixing_batch_shape):
        mixing_probabilities = jnp.repeat(
            jnp.expand_dims(mixing_probabilities, i), s, axis=i
        )
    assert jnp.allclose(mixing_probabilities.sum(axis=-1), 1.0)
    mixing_distribution = dist.Categorical(probs=mixing_probabilities)
    if same_family:
        component_distribution = jax_dist_getter((*component_batch_shape, nb_mixtures))
    else:
        component_distribution = [
            jax_dist_getter(component_batch_shape) for _ in range(nb_mixtures)
        ]
    _test_mixture(mixing_distribution, component_distribution)