def _partial_fn(fn, *args, **kwargs):
  """Partial apply *args and **kwargs after the shape, dtype, key args."""

  # Dummy value here is so that we can use lax.switch directly.
  def f(shape, dtype, key, _):
    return fn(shape, dtype, key, *args, **kwargs)

  return f