def KaimingUniformInitializer(out_dim=-1, in_dim=-2, param=0.):
  """Returns an initializer for random uniform Kaiming-scaled coefficients."""
  return ScaledInitializer(
      out_dim, in_dim, 2.0 / jnp.sqrt(1 + param**2), 'fan_in', 'uniform')
Returns an initializer for random uniform Kaiming-scaled coefficients.
trax/layers/initializers.py