Help to speed up code

[JohannesEsslinger]

Hello,
I'm having trouble properly speeding up my code. I simply want to draw samples from a Gaussian Mixture Model.

DIM=2
    
    def sample_from_component_fn(gmm_state, index, num_samples, seed):
        return jnp.transpose(jnp.expand_dims(gmm_state.means[index], axis=-1)
                             + gmm_state.chol_covs[index] @ jax.random.normal(key=seed, shape=(DIM, num_samples)))
    
    
    def sample_from_components_no_shuffle(gmm_state, samples_per_component: chex.Array,
                                          seed: chex.PRNGKey) -> [chex.Array, chex.Array]:
        samples = []

        for i in range(gmm_state.num_components):
            if samples_per_component[i] == 0:
                continue
            samples.append(sample_from_component_fn(gmm_state, i, samples_per_component[i], seed))
            seed, subseed = jax.random.split(seed)

        samples = jnp.vstack(samples)

        return samples

I cannot simply vmap (same for scan) over my number of components, as samples_per_component is different for each component and therefore the resulting shape, e.g.,

jax.vmap(sample_from_component_fn, in_axes=(None,0, 0, 0))(gmm_state, jnp.arange(
            gmm_state.num_components), samples_per_component, jax.random.split(seed, gmm_state.num_components))

I would really appreciate any help!

Regards

[jakevdp]

Hi, thanks for the question! Unfortunately, this represents somewhat of a worst case when it comes to efficient computation with JAX: you're looping over dispatches of many small functions, each of which operates on differently-sized arrays. As you discovered, vmap and scan won't help in this case. Further, jit won't be much help either, because it will re-compile the kernel each time it encounters a differently-sized input.

A common strategy in a case like this is to re-express each step in terms of operations over uniformly-shaped inputs (with padding). This looks like it may be possible in your case, because you're essentially just sampling from the input arrays at each step. You'd have to modify how gmm_state stores its state, to use constant-sized padded representations.

If rewriting in terms of padded operations is not possible, then your current approach is probably the best available.

Finally, side-note: the way you are splitting your seeds might lead to corrolation between the steps, because you are using seed twice in each step (and ignoring subseed). I'd instead use fold_in like this:

        for i in range(gmm_state.num_components):
            if samples_per_component[i] == 0:
                continue
            seed_i = jax.random.fold_in(seed, i)
            samples.append(sample_from_component_fn(gmm_state, i, samples_per_component[i], seed_i))

[JohannesEsslinger]

Thanks for the quick answer and for the hint with the seed!