Ensemble expected usage

[johnviljoen]

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.

    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?

[lockwo]

This is basically because of a partial commitment to a design decision I made. Originally I wanted to commit to exclusively handling batch dimensions via vmap (and here batch dimension is anything beyond the distribution state, so for Bernoulli, if you have an image of (28, 28) then it's 784 batch dimension since each distribution is a scalar). All code (I think) should be written to support vmaping based approaches. However, there are cases in which it may work without vmap-ing, but this is really more of a happy little accident (although admittedly, I didn't add errors/explicit warnings if there is some shape difference between the theoretical event shape and the detected shape). I say partial commitment, because I don't force things to be the theoretical event shape (like the bernoulli/Normal ones), so sometimes you can get away with implicitly batching things (like sampling in your example). Another reason I didn't fully restrict this was because in small cases, splitting many keys and generating things from independent keys can be quite slow relative to from one key (mostly on very sequential tasks like diffrax/SDEs).

The second two errors actually caused me a fair amount of headache early on as well, so maybe I should add more explicit notes about it.

TL;DR: those are expected errors and batching should be handled via vmaping.