Improve Error Message: Jitting linen.Module

Example code:

model = Model(1)

#@jax.jit
def eval_step(model, params, batch):
  logits = model.apply({'params': params}, batch['X'])
  return compute_metrics(logits, batch['y'])

def eval_model(model, params, test_ds):
  metrics = eval_step(model, params, test_ds)
  metrics = jax.device_get(metrics)
  summary = jax.tree_map(lambda x: x.item(), metrics)
  return summary['loss'], summary['accuracy']

print(eval_model(model, params, test_ds))

Throws the following general JAX error:

TypeError: Argument 'Model(
    # attributes
    features = 3
)' of type <class '__main__.Model'> is not a valid JAX type

Modules in Linen aren’t “pytypes” thus they can’t be flattened/unflattened as needed when entering and exiting JAX transformations. The common pattern is to use static_argnums to jit which is equivalent to having the module instance behave like a constant inside the transformed function

We should consider registering Modules with pytrees just to throw an error explaining this.

[jheek]

Normally you would use jax.jit(jitted_step_bad, static_argnums=(0,)) so the model is a static argument. You cannot pass a model as a normal arg because it’s not a (pytree of) JAX arrays.

[melissatan]

Was caught up with other responsibilities for some time, picking this back up now.

I’ve made the changes @jheek suggested, in the 3 tree_map() calls in get_module_scopes(). But this appears to be causing a test failure in several tests that use nn.jit() in linen_transforms_test, e.g.:

def test_compact_aliasing_collision(self):
    class Foo(nn.Module):
      m1: nn.Module
      m2: nn.Module
      @nn.compact
      def __call__(self, x):
        x = self.m2(self.m1(x))
        return x
    class Bar(nn.Module):
      @nn.compact
      def __call__(self, x):
        dense = nn.Dense(2)
        x = nn.jit(Foo)(dense, dense)(x)  # <-- fails here, and outputs the jit pytree error that I'm defining in the PR #2270 .
        return x
    k = random.PRNGKey(0)
    x = jnp.zeros((2, 2))
    _ = Bar().init(k, x)

I tried specifying static_argnums for the nn.jit() call in the test above, but then got another failure about “Non-hashable static arguments are not supported.”.

I’m not sufficiently familiar with the internals of linen_transforms. Is there another workaround that prevents the failure in module_lifecycle.rst (https://github.com/google/flax/issues/853#issuecomment-1113396067)?