Common Initializers Does Not Work with Bias

[daskol]

It is expected that bias vector can be initialized not only with zeros but the following code fails.

import jax, jax.numpy as jnp
import flax.linen as nn

layer = nn.Conv(features=32,
                kernel_size=(3, 3),
                use_bias=True,
                bias_init=nn.initializers.lecun_normal())

jax.jit(layer.init)(jax.random.PRNGKey(42), jnp.empty((5, 28, 28, 1)))

# File jax/core.py:1969, in NamedShape.__getitem__(self, idx)
#    1967 try:
#    1968   idx = operator.index(idx)
# -> 1969   return self.__positional[idx]
#    1970 except TypeError:
#    1971   pass
#
# IndexError: tuple index out of range

It seems that the issues #1386 is related and maybe #2002 is related as well.

[marcvanzee]

Assigning this to @zaxtax who is looking into improving our initializer documentaiton.

[daskol]

Thanks. It is not only about documentation. Is is about universality and interfacing of common initialize routines as well.

[marcvanzee]

Yes you are right!

[zaxtax]

@daskol I agree! Part of making good documentation is writing code that's easy to explain :)

[cgarciae]

The problem is that all initializes that are implemented on top of variance_scaling (like lecun_normal) require the shape to have at least 2 dimensions, so you can't pass these to bias_init. This should be documented, also maybe we should improve the error message.

[daskol]

Agreed. Documentation should be updated.

From my perspective, variance_scaling initializer can be extrapolated on bias term. Bias term can be considered as a edge case of variance_scaling-based initializer.