How to efficiently map a scalar-valued function over arrays with different leading axis shapes, but same remaining axes shapes?

[cisprague]

Suppose you have a dataset of particle configurations, where each configuration is an array of positions of shape (n, d) where n ∈ {0, ..., n_max} is finite and d is constant. Consider you have a single-sample loss function loss_sample that computes a scalar loss for each configuration. Since n varies per example, we cannot simply vmap loss_sample over the whole batch of configurations. But, since n_max is finite, and most likely much less than the number of examples (n_max << n_configs) in the dataset, it might be reasonable to jit the loss_sample function for each example (with JAX maintaining a cache of compiles). Hence, loss_sample will be compiled at most n_max times. But, how can we then create a loss_batch that efficiently maps these jited loss_sample functions over a batch of arrays with different n? If it is possible to do this efficiently, would it be a viable alternative to padding/masking? See below.

def gen_batch(d, n_max, n_configs, key):
    ns = jax.random.randint(key, (n_configs,), 1, n_max)
    configs = [jax.random.normal(key, (n, d)) for n in ns]
    return configs

@jax.jit
def loss_sample(config):
    print('Compile!')
    return jnp.sum(jnp.linalg.norm(config, axis=1))

def loss_batch(configs):
    return sum([loss_sample(config) for config in configs]) # NOTE: works
    # return jnp.sum(jax.vmap(loss_sample)(configs)) NOTE: doesn't work

configs = gen_batch(3, 10, 100, jax.random.PRNGKey(0))
loss_batch(configs)

>>> configs = gen_batch(3, 10, 100, jax.random.PRNGKey(0))
>>> loss_batch(configs)
Compile!
Compile!
Compile!
Compile!
Compile!
Compile!
Compile!
Compile!
Compile!
Array(918.0347, dtype=float32)
# less than 100 compiles

[jakevdp]

You might try using lax.map with a batch_size smaller than your smallest set of points. This would end up performing multiple vmapped batches sequentially.