Cannot use `flax.linen.initializers.zeros` or `ones` with `Module.variable`
- Dominant language
- Jupyter Notebook
- Stars
- 7.3k
- Forks
- 833
- Avg merge
- 5h 11m
- Merged PRs (30d)
- 5
Description
When attempting to use the `flax.linen.initializers.zeros` or `flax.linen.initializers.ones` initializers with `Module.variable`, I get:
> TypeError: zeros() missing 1 required positional argument: 'shape'
Here is a minimal example:
```
from flax import linen as nn
import jax
from jax import numpy as jnp
class Example(nn.Module):
@nn.compact
def __call__(self, x):
sample_variable = self.variable(
'param', 'mean', nn.initializers.zeros, (5, ))
return
rng = jax.random.PRNGKey(0)
x = jnp.array([1,2,3])
Example().init(rng, x)
```
This is the correct way to use `Module.variable`, right? This is how it is used in the [linen upgrade guide](https://docs.google.com/document/d/1hYavTVPaKVVe9Be8pCB7yW7r6dDv3RALVNit8NZca4c/edit?pli=1#heading=h.l18t12ge3ai4) which I am using to port some pre-linen flax code.
Replacing `nn.initializers.zeros` with `functools.partial(nn.initializers.zeros, shape=(5,)` fixes the problem, but the call to `self.variable` should already be passing the shape to `zeros`, right?
### System information
- OS Ubuntu 20.04.5 LTS on Windows 10 x86_64
- Flax 0.6.10, jax 0.4.11, jaxlib 0.4.11
- Python version: 3.8.16
- GPU/TPU model and memory: CPU
Contributor guide
Assessment
This issue has not been assessed yet.