Help with lifted vmap

[marcosrdac]

I've recently created a module that simulates wave propagation. It is written as below:

from typing import Any, Iterable, Tuple
from jax import numpy as jnp, random
from flax import linen as nn

class WERNN(nn.Module):
    v0: jnp.array
    dz: float
    dx: float
    dt: float

    def makesimulation(self, source: Any, nt: int):
        '''
        Returns acoustic wave simulation results given a source and a number of time 
        steps to be calculated (nt).
        '''
        return jnp.empty((self.v0.shape[0], nt))  # just for simplification

    @nn.compact
    def __call__(self, source: Any, nt: int):
        self.variable('variables', 'something', lambda: 0)
        return self.makesimulation(source, nt)

This module returns a 2D array. As I need to use the result of many simulations stored in a 3D array, I thought of doing a lifted vmap on the WERNN module. Here is my trial for this simplified model:

class BatchWERNN(nn.Module):
    v0: jnp.array
    dz: float
    dx: float
    dt: float

    @nn.compact
    def __call__(self, sources: Tuple[Any], nt: int):
        WERNNVmap = nn.vmap(WERNN, in_axes=0, out_axes=0,
                            variable_axes={'params': None, 'variables': None},
                            split_rngs={'params': False, 'variables': False})

        return WERNNVmap(v0=self.v0, dz=self.dz, dx=self.dx, dt=self.dt)(sources, nt)

In my example nt is meant to be an integer that is used at each WERNN call. sources is a tuple of values that should each be used at each simulation done (WERNN call). But here is the error I get:

    raise ValueError(f"{name} got arg {i} of rank {len(shape)} but axis to be "
ValueError: vmap got arg 1 of rank 0 but axis to be mapped 0. The tree of ranks is:
(((({},), ({},)), (({},), ({},)), (1, 0)), {})

I still cannot understand what is getting in my way though. Do you? I really appreciate your attention. Thank you for trying to help me!

[jheek]

I think you need to broadcast nt because you reuse the same value in each call. So in_axes=(0, None)

[marcosrdac]

@jheek, thanks for your attention. I remember having already tried adding this not to broadcast nt, and it did not work in my full problem. But you are right, it does work in the one I provided here. I believe I found the inconsistency now. What happens is that sources is actually a tuple of arrays (you should be able to reproduce my problem adding next code to the previous ones):

    rng = random.PRNGKey(0)

    # model definition
    shape = nz, nx = 4, 5
    v0 = jnp.ones(shape)
    model = BatchWERNN(v0, dz=100, dx=100, dt=.008)

    # sources definition
    source1 = jnp.asarray([0, 0, 0]), jnp.asarray([1., 2.])
    source2 = jnp.asarray([0, 1, 0]), jnp.asarray([1., 2., 3., 4., 5.])
    sources = source1, source2

    # running model
    params = model.init(rng, sources, 3)  # gives us an error
    result = model.apply(params, sources, 3)

    print(result.shape)

I need to pass 2 arrays in each element inside sources, but while the first array has shape (3,) (an array of integers meaning a position), the second (the source signal) has shape (n,) (where n is variable for each source, as the span of each signal in time might differ). I believe the problem lies in there. My belief is based on the fact that using:

    source1 = jnp.asarray([0, 0, 0]), jnp.asarray([1., 2., 3.])
    source2 = jnp.asarray([0, 1, 0]), jnp.asarray([1., 2., 3.])

works nicely. Is there a way to pass such sources of variable shapes to a lifted vmap? Else, which approach do you think I should try in order to correct my code?

(I don't know if I should open a new discussion now, as my problem turned out to be lifted vmapping with variable shapes)

[jheek]

Vectorize requires that the batched dimension is the same for each input.
My guess is that you actually want to vectorize over the element in the tuple.
But jax just sees each element of the tuple as another input argument.
So you should probably make an array of shape (2, 3) of source1 and shape (2, n)
JAX however does not (yet) support ragged batching so n has to be fixed for each input.

[marcosrdac]

I see what you mean, @jheek. I confirm that you are right. I'll settle and pad the arrays so that they are all the same size. Thank you for your tips!