def AtariConvInit(kernel_shape, rng, dtype=jnp.float32):
  """The standard init for Conv laters and Atari."""
  filter_height, filter_width, fan_in, _ = kernel_shape
  std = 1 / jnp.sqrt(fan_in * filter_height * filter_width)
  return random.uniform(rng, kernel_shape, dtype, minval=-std, maxval=std)
The standard init for Conv laters and Atari.
trax/layers/initializers.py