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])
    ...

Please provide the full traceback of any errors.

  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..

[rlouf]

We do not intend to support the neural net code as part of the core Aesara library (#674). The code may however be moved to a separate library at a later point in time.