How to dispatch to different instances of the same Module?

[marcvanzee]

Q

[marcvanzee]

Answer by @levskaya:

Dynamically dispatching can have complicated performance behavior in the XLA setting (if you are rnning on TPU or GPU).
Why not just computing all results and indexing?

class Foo(nn.Module):
  @nn.compact
  def __call__(self, x, idx):
    denses = [nn.Dense(2) for i in range(3)]
    ys = [denses[i](x) for i in range(3)]
    return ys[idx]