Use vmap to apply the same model simultaneously to batch of inputs

[imi-hub]

Hello!

I have a list of models which I would like to apply the same batch of inputs.

Instead of using a for loop to go through the models in the list I was wondering if it is possible to parallelize the application of models onto the inputs with vmap.

Similar issue:
#673

I share a simple example below:

import flax.linen as nn
import jax
import jax.numpy as jnp
from jax import random


class Model(nn.Module):
    alpha: int = 1
    @nn.compact
    def __call__(self, x):
        y = nn.Dense(features=self.alpha*x.shape[0])(x)
        y = nn.relu(y)
        return jnp.sum(y, axis=-1)
    
models = [Model(alpha=1) for i in range(3)]



# Define random input tensor and params of model
batch_size = 2
input_shape = (4,)
x = random.normal(random.PRNGKey(0), (batch_size,) + input_shape)
params = model.init(random.PRNGKey(0), x)

Below is the attempt with a for loop, which works fine, but it is slow (for more complex models with larger amount of parameters...)

# LOOP through the models list (3 models)
model_out = []
for model in models:
    model_out.append(model.apply(params, x))
print('\nLoop: model output\n', model_out)

1st attempt

Below I share my first attempt to vmap the model of models to the inputs:

# Try to vmap over the models list to apply the neural networks in parallel to the input x

# 1st attempt
def model_fct(model, n):
    return model(n)

vmap_out1 = jax.vmap(model_fct, in_axes=(0,None))(models,x) #jax.vmap(model_apply, in_axes=(0,None))(models, x)

print('\n1-vmap: model output\n', vmap_out1)

which gives me the following error:

---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
File ~/netket_env/lib/python3.9/site-packages/jax/_src/api.py:1664, in _mapped_axis_size.<locals>._get_axis_size(name, shape, axis)
   1663 try:
-> 1664   return shape[axis]
   1665 except (IndexError, TypeError) as e:

IndexError: tuple index out of range

The above exception was the direct cause of the following exception:

ValueError                                Traceback (most recent call last)
Cell In [238], line 42
     39 def model_fct(model, n):
     40     return model(n)
---> 42 vmap_out1 = jax.vmap(model_fct, in_axes=(0,None))(models,x) #jax.vmap(model_apply, in_axes=(0,None))(models, x)
     44 print('\n1-vmap: model output\n', vmap_out1)
     47 # 2nd attempt

    [... skipping hidden 6 frame]

File ~/netket_env/lib/python3.9/site-packages/jax/_src/api.py:1667, in _mapped_axis_size.<locals>._get_axis_size(name, shape, axis)
   1665 except (IndexError, TypeError) as e:
   1666   min_rank = axis + 1 if axis >= 0 else -axis
-> 1667   raise ValueError(
   1668       f"{name} was requested to map its argument along axis {axis}, "
   1669       f"which implies that its rank should be at least {min_rank}, "
   1670       f"but is only {len(shape)} (its shape is {shape})") from e

ValueError: vmap was requested to map its argument along axis 0, which implies that its rank should be at least 1, but is only 0 (its shape is ())

2nd attempt

And below is a second attempt:

# 2nd attempt
vmap_out = jax.vmap(lambda i, x: jax.lax.switch(i, models, x), in_axes=(0,None)) #jax.vmap(model_apply, in_axes=(0,None))(models, x)
print('\nvmap: model output\n', vmap_out(jnp.arange(len(models)),x))

with following error:

CallCompactUnboundModuleError             Traceback (most recent call last)
Cell In [239], line 49
     47 # 2nd attempt
     48 vmap_out = jax.vmap(lambda i, x: jax.lax.switch(i, models, x), in_axes=(0,None)) #jax.vmap(model_apply, in_axes=(0,None))(models, x)
---> 49 print('\nvmap: model output\n', vmap_out(jnp.arange(len(models)),x))

    [... skipping hidden 3 frame]

Cell In [239], line 48, in <lambda>(i, x)
     38 '''# 1st attempt
     39 def model_fct(model, n):
     40     return model(n)
   (...)
     45 
     46 '''
     47 # 2nd attempt
---> 48 vmap_out = jax.vmap(lambda i, x: jax.lax.switch(i, models, x), in_axes=(0,None)) #jax.vmap(model_apply, in_axes=(0,None))(models, x)
     49 print('\nvmap: model output\n', vmap_out(jnp.arange(len(models)),x))

    [... skipping hidden 14 frame]

File ~/netket_env/lib/python3.9/site-packages/flax/linen/module.py:818, in Module._call_wrapped_method(self, fun, args, kwargs)
    816 if is_compact_method:
    817   if self.scope is None:
--> 818     raise errors.CallCompactUnboundModuleError()
    819   is_recurrent = self._state.in_compact_method
    820   self._state.in_compact_method = True

CallCompactUnboundModuleError: Can't call compact methods on unbound modules (https://flax.readthedocs.io/en/latest/api_reference/flax.errors.html#flax.errors.CallCompactUnboundModuleError)

[jakevdp]

Hi - thanks for the question! Since this is pretty specific to flax, you may have more luck at https://github.com/google/flax/discussions/

The answer to why this is not working is that vmap requires array inputs, not list inputs. Your list of models is an "array-of-structs" data structure, whereas vmap works with a "struct-of-arrays" data structure. In other words, you'd need to somehow batch the model parameters rather than having a list of individual models. As for how to do that with flax models, I'm not sure, but the flax discussions are monitored by people who might know how to approach it.

[imi-hub]

thank you!

[cgarciae]

Hey @imi-hub, if all models have the same form you can use nn.vmap to create an ensamble. If you want all instance to receive the same data but have different weights set in_axes=None, variable_axes={'params': 0}, split_rngs={'params': True}, and axis_size equal to the number of models.

[imi-hub]

thank you!