def _unshaped_initializer(fn,
                          **kwargs):
  """Build a flax-style initializer that ignores shapes.

  Args:
    fn: Callable that takes a random number generator and some keyword
      arguments, and returns an initial parameter value.
    **kwargs: Arguments to pass to fn.

  Returns:
    Callable that takes a PRNGKey and an (unused) shape, and returns the value
    returned by fn.
  """

  def _initialize(rng_key, unused_shape):
    return fn(rng_key, **kwargs)

  return _initialize