google / google/flax

Cannot use `flax.linen.initializers.zeros` or `ones` with `Module.variable`

Open
#3,140 1 comment 0 reactions 0 assignees View on GitHub
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

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.