Walk the theta and observation blocks of an MDN's i.i.d. density, calling collect(theta_idx, obs_idx, lp) with the (n_theta_chunk, n_obs_chunk) tensor of log-densities

The MLP maps theta to the mixture parameters and never sees x, so within one theta block the forward pass and the Cholesky assembly run once and only the quadratic form is chunked over observations. theta itself is blocked first (mdn_theta_chunk_size()), so max_batch bounds mdn_mixture()'s output the same way cross_iid() bounds a flow's forward pass, rather than only chunking the observation side (#240).

Usage

mdn_iid_blocks(de, xt, theta, max_batch, collect)