Multiple initializations. Is this a bug?

I've run into a strange behavior, and I'm unsure if its a bug or if I'm doing something wrong, and I wasn't able to find any clarification in the docs: it seems that the weights of my network get initialized every time `apply`is called, instead of just the first time, as seen in the example below where I create a custom initialization function. This seems to be that it would be inefficient, so I'm wondering if I'm missing something.

```python
class MyModule(nn.Module):

def setup(self):

def my_bias_init(rng, shape, dtype):
print('bias init')
return random.uniform(rng, shape, dtype=dtype, minval=0, maxval=2)

self.a_layer = nn.Dense(10, bias_init=my_bias_init)
```

**System information**
```

Name: flax
Version: 0.7.4
Summary: Flax: A neural network library for JAX designed for flexibility
Home-page:
Author:
Author-email: Flax team
License:
Location: /2tb_nvme/hpinkard_waller/mambaforge/envs/phenotypes/lib/python3.10/site-packages
Requires: jax, msgpack, numpy, optax, orbax-checkpoint, PyYAML, rich, tensorstore, typing-extensions
Required-by:
---
Name: jax
Version: 0.4.18
Summary: Differentiate, compile, and transform Numpy code.
Home-page: https://github.com/google/jax
Author: JAX team
Author-email: jax-dev@google.com
License: Apache-2.0
Location: /2tb_nvme/hpinkard_waller/mambaforge/envs/phenotypes/lib/python3.10/site-packages
Requires: ml-dtypes, numpy, opt-einsum, scipy
Required-by: chex, flax, optax, orbax-checkpoint
---
Name: jaxlib
Version: 0.4.18+cuda12.cudnn89
Summary: XLA library for JAX
Home-page: https://github.com/google/jax
Author: JAX team
Author-email: jax-dev@google.com
License: Apache-2.0
Location: /2tb_nvme/hpinkard_waller/mambaforge/envs/phenotypes/lib/python3.10/site-packages
Requires: ml-dtypes, numpy, scipy
Required-by: chex, optax, orbax-checkpoint

```