quantax.nn.apply_he_normal

quantax.nn.apply_he_normal(key: Key, net: Linear | Conv) → Linear | Conv

Apply the He normal initializer to the weights of the layer. The bias is initialized to 0.

Parameters:

• key – The random key used in JAX for initializing parameters.

• net – The net to apply the initializer.

Returns:

The net with properly initialized parameters.

Note

This function can only be applied to Linear or Conv layers in Equinox.

Note

The input net is not modified.