google / google/flax

Properly bind all submodules

Open
#2,744 2 comments 0 reactions 0 assignees View on GitHub
Priority: P2 - no schedule
Dominant language
Jupyter Notebook
Stars
7.3k
Forks
833
Avg merge
5h 11m
Merged PRs (30d)
5

Description

Currently submodules are lazily bounded when `__getattr__` that calls `_try_setup` , however, this only works for submodules defined in `setup`, submodules that are passed as constructor arguments don't trigger `__getattr__` (since their fields always exist) and are never bounded e.g:

```python
import flax.linen as nn
import jax
import jax.numpy as jnp

module = nn.Sequential([
nn.Dense(8),
nn.relu,
nn.Dense(4),
])
x = jnp.ones((1, 5))
variables = module.init(jax.random.PRNGKey(0), x)
bound_module = module.bind(variables)

print(bound_module.scope) #
print(bound_module.layers[0].scope) # None <== !!!
```
This adds some rough edge cases when using `bind` as some submodules can be unavailable.

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.