Parameter initialization
Specialized parameter initializers, extending jax.nn.initializers.
For some initialization schemes, the bias depends on the number of input features, which cannot be determined from the shape of the bias array (see, e.g, Hoedt normal initializer). To handle such cases klax provides a custom klax.Initializer protocol that generalizes upon the JAX API, while ensuring compatibility with all jax.nn.initializers.
klax.Initializer
¤
Protocol for initializers, generalizing jax.nn.initializers.Initializer.
Some advanced initialization schemes initialize the bias
depending on the number of input features (fan_in). However,
from the bias shape alone fan_in cannot be computed. This
protocol specifies an initializer that is supplied with
fan_in explicitly, enabling advanced bias initialization.
__call__(key, shape, fan_in, dtype=float64)
¤
klax.hoedt_normal(in_axis=-2, dtype=float64)
¤
Build a Hoedt normal initializer (for positivity constrained weights).
A Hoedt normal initializer is designed
for weights with positivity constraint, such as in klax.nn.FICNN. It
samples weights according to a log-normal distribution derived by studying
signal propagation through layers with non-negative weights.
Tip
This initiailzation should be paired with the klax.hoedt_bias
initializer for biases of constrained layers.
| PARAMETER | DESCRIPTION |
|---|---|
in_axis
|
Axis of the input dimension in the weights array.
TYPE:
|
dtype
|
The dtype of the weights.
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
Initializer
|
A |
klax.hoedt_bias()
¤
Build a Hoedt bias initializer (for layers with positivity constrained weights).
A Hoedt bias initializer is designed
for the unconstrained biases in linear layers where the weights are constrained
to be positive, such as in klax.nn.FICNN. It intializes biases to vectors of
a constant value, computed from the number of input features (fan_in).
Tip
This initialization should be paired with the klax.hoedt_normal initializer
for the positivity constraint weights.
| RETURNS | DESCRIPTION |
|---|---|
Initializer
|
An |