Add reinterpreted_batch_ndims to Independent

[gvcallen]

Adds reinterpreted_batch_ndims to Independent, to support distributions that have been batched via eqx.filter_vmap (e.g., MultivariateNormalTri) that don't support handle batched arrays natively.

Although distreqx has expressly dropped the concept of batch dimensions throughout the codebase (which I strongly agree was the right choice), I do believe Independent still requires a reinterpreted_batch_ndims options, with motivation in the example below.

Without this PR, passing a vmapped distribution that doesn't natively support broadcasing to Independent crashes on methods like log_prob, because the inner computations weren't being appropriately mapped. This change allows users to explicitly define the number of mapped axes, letting Independent correctly unroll the vmap layers and reduce the results to a single scalar. The previous behaviour is left unaffected with the default of reinterpreted_batch_ndims = 0.

locs = jnp.zeros((20, 10, 3))
scales_tri = jnp.stack([jnp.tri(3)]*20*10, axis=0).reshape(20, 10, 3, 3)
xs = jnp.ones((20, 10, 3))

# Create a batch of mvn's
mvns = eqx.filter_vmap(eqx.filter_vmap(dist.MultivariateNormalTri))(locs, scales_tri)

# Calling log_prob (expectedly) does not work because MultivariateNormalTri assumes k x k
# log_prob = mvn.log_prob(locs) # error

# We can manually vmap the computation, but then we get a vector of log-probs.
# We could add a sum, but it is better to encapsulate this in a single distribution
def simple_log_prob(d, x): return d.log_prob(x)
log_probs = eqx.filter_vmap(eqx.filter_vmap(simple_log_prob))(mvns, xs)
assert log_probs.shape == (20, 10,)

# If we attempt to use "Independent" to reinterpret the batch dims, we get a crash,
# because the computation dimensions are not being reinterpreted
reinterp_mvn = dist.Independent(mvns)
# reinterp_mvn.log_prob(xs) # error

# With the propose changes, we can successfully encapsulate the independent MVNs.
# We could also pass e.g. 1 here and manually vmap the last dim if desired.
reinterp_mvn = dist.Independent(mvns, reinterpreted_batch_ndims=2)
log_prob = reinterp_mvn.log_prob(xs)
assert jnp.isscalar(log_prob)

[lockwo]

Does this note render correctly? I think we can make nice note interface. Here is an AI summary of the interfaces

Solid block (!!!) — always visible:                                                                                                
  !!! note        
                                                                                                                                     
      Body must be indented 4 spaces and separated from the
      `!!!` line by a blank line.                                                                                                    
                                 
  Collapsible block (???) — collapsed by default; user clicks to open:                                                               
  ??? note                                                                                                                           
          
      Same indentation rules.                                                                                                        
                                                                                                                                     
  Open-by-default collapsible (???+):
  ???+ warning                                                                                                                       
                  
      Starts expanded but is still collapsible.                                                                                      
                                               
  Custom title — quoted string after the type:                                                                                       
  !!! warning "Bijectors are applied in reverse order"                                                                               
                                                      
      Given a sequence `[f, g]`, the `Chain` bijector computes `f(g(x))`... 

[lockwo]

I'm not sure we should have a specific flag that goes to the reduce helper, these means it should reduce over all dimensions basically, should it be instead a required argument (that then becomes all axis rather than just 0)? Just to simplify code?

[lockwo]

this patterns comes up several times, I think if we make a vmap helper it could reduce LoC? e.g.

  def _vmap_method(self, fn,):                                                                                                 
      for _ in range(self.reinterpreted_batch_ndims):                                                                              
          fn = eqx.filter_vmap(fn)                                                                                                   
      return fn(self.distribution,)


  def _vmap_and_sum(self, fn,):                                                                                                
      out = self._vmap_method(fn,)                                                                                           
      return jnp.sum(out, axis=tuple(range(self.reinterpreted_batch_ndims)))     

then each of these can just call vmap and sum

[lockwo]

what is fmt skip for here?

[lockwo]

if these don't match this has an asymmetric issue? should we enforce this eto be the same

[lockwo]

un-needed comments

[lockwo]

personally, I don't think we need to have much backwards compatible, we can break things if the new API is better.

[lockwo]

we should check across multiple distributions, e.g. bernoulli/normal/etc to make sure nothing unexpected happens