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