How do I scan / map / vmap over a function of variable sized inputs?

[ahwillia]

Doing this:

def foo(x):
    return x.sum()
jax.lax.map(foo, [jnp.array([1, 2]), jnp.array([1, 2, 3])])

Causes ValueError: scan got values with different leading axis sizes: 2, 3.... Is there a standard way to work around this?

This kind of thing can come up if you have batched data of different sizes, for example.

[YouJiacheng]

The standard way is padding data to the same size.
For better efficiency, you can bucket your data by size, sample data by sampling a bucket then sampling a batch of data in this bucket instead of directly sampling a batch of data.

[YouJiacheng]

Hmmm, you iterate over pytrees? To avoid XY problem, I suggest you provide more information.

[ahwillia]

I'm computing the log-posterior of a hierarchical Bayesian model. The data comes in groups / batches with varying numbers of observations. There are local parameters to each group that have different shapes, since the number of observations is different in each group. It seems natural to iterate over something like:

((param1_batch1, param2_batch1), data_batch1), ((param1_batch2, param2_batch2), data_batch2), ...

And for each call a function that returns the log probability

def logprob(args):
    (param_1, param_2), data = args    # Unpack params and data
    return ... # some computation that depends on params, data

This is a simple example with two params per group, but in general there could be lots of params.

I have another problem which is how to share global information among each batch efficiently without creating impure functions, but that is a separate issue.

[YouJiacheng]

Do param1_batch1, param2_batch1 and data_batch1 have different number of observations,
OR batch1 and batch2 have different number of observation?
And you want to process batch1 and batch2 at the same time(why?).

[ahwillia]

At each iteration I want to hand off a single batch, ((param1_batch1, param2_batch1), data_batch1). The shape of data_batch1 includes a dimension num_observations that varies across batches. In addition, some (maybe not all) of the parameters also have dimensions that depend on the number of observations. I'm trying to keep this high-level because I think this is a reasonably common situation.

(I was able to refactor my problem last night to get something working, but I threw away the functional programming aesthetic that is typical in using jax. I used a big sparse matrix multiply and other vectorized ops to compute the loss over all batches at once.)

[YouJiacheng]

For such situation, I think we typically just use python for-loop to iterate over batches. Note that will trigger compiling for each different batch size if you use jit. So you can bucket your batches by size, and padding them to maximum size of the bucket, i.e. (0, 4] padded to 4, (4, 8] padded to 8.
You can use tree_map to help you pad the data:

data = ((param1_batch1, param2_batch1), data_batch1)
batch_axis = ((None, 1), 0) # tree-prefix of data-pytree-structure , None stand for not depending on number of observations
def pad(arr, axis):
    if axis is None:
        return arr
    # pad your array ...
padded_data = jax.tree_util.tree_map(lambda axis, data_leaf: pad(data_leaf, axis), batch_axis, data)