`Module.init` does not contain KV-pair of modules without parameters
- 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
Assessment
This issue has not been assessed yet.