Initialization of `flax.linen` model parameters depends on member variable names
- Dominant language
- Jupyter Notebook
- Stars
- 7.3k
- Forks
- 833
- Avg merge
- 5h 11m
- Merged PRs (30d)
- 5
Description
Provide as much information as possible. At least, this should include a description of your issue and steps to reproduce the problem. If possible also provide a summary of what steps or workarounds you have already tried.
### System information
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Linux Ubuntu 22.04
- Flax, jax, jaxlib versions (obtain with `pip show flax jax jaxlib`: (flax: 0.8.5, Jax: 0.4.28, JAXlib: 0.4.28)
- Python version: 3.11
### Problem you have encountered:
It appears that the initialization of model parameters in `flax.linen` depends on the name of the submodule in the parent model, rather than solely on the provided `PRNGKey`. When two models use the same module but with different submodule names, the parameters are initialized differently.
### What you expected to happen:
I expected the parameter initialization to depend only on the provided PRNGkey (`jax.random.PRNGKey(0)`) and not on the names of submodule in the parent model.
### Logs, error messages, etc:
None, but the observed output when running the example code is:
```
{'1': {'params': {'net': {'Dense_0': {'bias': Array(0., dtype=float32),
'kernel': Array(2.1496024, dtype=float32)}}}},
'2': {'params': {'net2': {'Dense_0': {'bias': Array(0., dtype=float32),
'kernel': Array(1.4864768, dtype=float32)}}}},
'3': {'params': {'net': {'Dense_0': {'bias': Array(0., dtype=float32),
'kernel': Array(2.1496024, dtype=float32)}}}}}
```
### Steps to reproduce:
```python
import jax, jax.numpy as jp
import flax.linen as nn
import optax
from pprint import pp
class Tmp(nn.Module):
@nn.compact
def __call__(self, x):
return nn.Dense(2)(x)
class Network(nn.Module):
net: Tmp
def __call__(self, x):
return self.net(x)
class Network2(nn.Module):
net2: Tmp
def __call__(self, x):
return self.net2(x)
class Network3(nn.Module):
net: Tmp
def __call__(self, x):
return self.net(x)
net = Tmp()
model1 = Network(net)
model2 = Network2(net)
model3 = Network3(net)
rng = jax.random.PRNGKey(0)
x = jp.zeros((1, 3))
params = {
"1": model1.init(rng, x),
"2": model2.init(rng, x),
"3": model3.init(rng, x),
}
pp(jax.tree_map(optax.global_norm, params))
```
Contributor guide
Assessment
This issue has not been assessed yet.