static_argnums argument to flax.linen.remat not working as expected

I am using `flax.linen.remat` on a module that has a `train` flag (used to check if the model is training). I'm using `static_argnums` on that flag, but am still getting a `ConcretizationTypeError` on model init.

Reproducer:

```python
import jax
import jax.numpy as jnp
import flax.linen as nn

class MLP(nn.Module):
@nn.compact
def __call__(self, x, train):
x = nn.BatchNorm(use_running_average=not train)(x)
x = nn.Dense(512)(x)
x = nn.relu(x)
x = nn.Dense(512)(x)
x = nn.relu(x)
x = nn.Dense(1)(x)
return x

# Works fine with this line commented out
MLP = nn.remat(MLP, static_argnums=(1,))

model = MLP()
rng_key = jax.random.PRNGKey(42)
variables = model.init(rng_key, input_example, True)
```

Traceback:

```bash
$ python foo.py
jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
File "/Users/dion/codes/supersede/foo.py", line 43, in
variables = model.init(rng_key, input_example, True)
jax.errors.ConcretizationTypeError: Attempted boolean conversion of traced array with shape bool[]..
The error occurred while tracing the function new_fun at /Users/dion/.virtualenvs/science/lib/python3.10/site-packages/jax/_src/ad_checkpoint.py:393 for checkpoint. This concrete value was not available in Python because it depends on the value of the argument dyn_args[2].
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerBoolConversionError

Consider using the `static_argnums` parameter for `jax.remat` or `jax.checkpoint`. See the `jax.checkpoint` docstring and its example involving `static_argnums`:
https://jax.readthedocs.io/en/latest/_autosummary/jax.checkpoint.html
```

Tested with `flax==0.8.4`.