Ensemble expected usage

Nobody has claimed this yet.

I am trying to integrate distreqx into some work I am currently doing and I had a question about expected usage.

I need a batch of MultivariateNormalDiag's so I have been creating ensembles with eqx.filter_vmap. Upon creation I can then .sample(key) the created vmapped distribution successfully (with a non-vmapped key), but other methods of the distributions class cannot work without first being integrated.

```python
eqx_pi = eqx.filter_vmap(dist.MultivariateNormalDiag)(
eqx_mean,
eqx_scale,
)
eqx_pi_log_prob = lambda d, a: d.log_prob(a)
eqx_pi_entropy = lambda d: d.entropy()
eqx_samp = eqx_pi.sample(key) # works
ent = eqx.filter_vmap(eqx_pi_entropy)(eqx_pi) # works
lp = eqx.filter_vmap(eqx_pi_log_prob)(eqx_pi, eqx_samp) # works

ent = eqx_pi.entropy() # evaluates, incorrect values, correct shape
lp = eqx_pi.log_prob(eqx_samp) # evaluates, incorrect values (all same), correct shape
```

Why can .sample work in this case (is it specific to my example)? Is this expected behaviour/am I using this incorrectly?

Start with MultivariateNormalDiag and the sample, entropy, and log_prob entry points shown in the issue, then inspect how eqx.filter_vmap wraps the distribution. Compare the direct method calls with the working vmapped lambdas and determine whether the behavior is expected or needs a library change. Done means documenting the expected usage or reproducing and fixing the incorrect values.

Written by the indexing model from the issue text.