Selecting a submodule from a list according to a tracer

[icysapphire]

I have a pool of convolutional layers and I would like to select which layer to use depending on the value of a tracer. The following is the sketch of my custom module:

def setup(self):
    self.pool_conv = [nn.Conv(features=3,
                                  kernel_size=(7, 7),
                                  kernel_init=get_init_fn(self.pool_init),
                                  bias_init=get_init_fn("zeros"),
                                  strides=(2, 2),
                                  padding=((3, 3), (3, 3)),
                                  use_bias=True) for _ in range(self.size)]


@nn.compact
def __call__(self, x, prompt_mask=None, task_id=-1, query=None, stop_gradients=False):
    ....
    initializing = self.is_mutable_collection('params')
    if initializing:
        for j in range(self.size):
            logging.info('making sure that conv%s is intiialized', j)
            tmp = self.pool_conv[j](x[0])

    @partial(jit, static_argnums=0)
    def get_conv_from_pool(conv_idx, sample):
        return self.pool_conv[conv_idx](sample)

    conv_out = get_conv_from_pool(conv_idx=jax.device_get(conv_idx).item(), sample=x[i])
    ...

The traceback:

  File "/clp/models/module.py", line 264, in __call__
    conv_out = get_conv_from_pool(conv_idx=jax.device_get(conv_idx).item(), sample=x[i])
  File "/miniconda3/envs/fo4/lib/python3.9/site-packages/jax/_src/device_array.py", line 226, in item
    return int(self)
jax._src.errors.ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(int32[])>with<DynamicJaxprTrace(level=0/1)>
The problem arose with the `int` function. If trying to convert the data type of a value, try using `x.astype(int)` or `jnp.array(x, int)` instead.

I am a newbie and I am trying to figure out a possible implementation considering JAX compilation model.. can you help me? It seems that static args do not help here..

[jakevdp]

Hi - the short answer is no, you cannot use a tracer to index into a list. But in the specific case of selecting from a list of functions to apply based on a traced value, jax.lax.switch is the tool you're looking for. In your case, its use would look something like:

# return self.pool_conv[conv_idx](sample)
return jax.lax.switch(conv_idx, self.pool_conv, sample)

There is one caveat: switch only works if the functions all return values with the same static properties (generally shape and dtype). If this is not the case, then what you're trying to do is fundamentally incompatible with JAX's requirement that within jit, array shapes be static, and not dependent on runtime data.