google / google/flax

`Module.init` does not contain KV-pair of modules without parameters

Open
#2,777 4 comments 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

### System information
- `Ubuntu 18.04.6 LTS`
- Flax, jax, jaxlib versions (obtain with `pip show flax jax jaxlib`:
```
Name: jax
Version: 0.3.25
---
Name: jaxlib
Version: 0.3.25+cuda11.cudnn805
---
Name: flax
Version: 0.6.3
```
- Python version: Python 3.8.16

### Problem you have encountered:
The dictionary of parameters generated at Module initialisation does not include keys of the modules that do not contain any parameter

### What you expected to happen:
That modules without parameters are represented as empty dictionaries, rather than being omitted completely.

### Steps to reproduce:
Whenever possible, please provide a *minimal example*. Please consider submitting it as a Colab link.
```python
import jax
import flax.linen as nn

class Identity(nn.Module):
def __call__(self, x):
return x

net = nn.Sequential([Identity(), nn.Dense(2)])

params = net.init(jax.random.PRNGKey(0), jax.numpy.ones((2, 2)))
print(params)
```

Current output:
```
# FrozenDict({
# params: {
# layers_1: {
# kernel: ...,
# bias: ...,
# },
# },
# })
#
```

Expected output:
```diff
# FrozenDict({
# params: {
+# layers_0: {},
# layers_1: {
# kernel: ...,
# bias: ...,
# },
# },
# })

```

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.