google / google/flax

Initialization of `flax.linen` model parameters depends on member variable names

Open
#4,367 1 comment 1 reaction 0 assignees View on GitHub
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.