Understanding un-normalized log probability in jax.random.categorical

[sh0416]

I am wondering whether the input of jax.random.categorical could be actually un-normalized log probability.

I think that the input would be normalized log probability.

jax.random.categorical leverages the Gumbel-max trick, used in various papers.

In this paper, they explain the trick using normalized log probability and I think it makes sense.

However, when we use un-normalized log probability, the scale of sampled value from Gumbel distribution remains unchanged, but the scale of un-normalized log probability would be changed, which means that the sampled result would be affected depending on the scale of un-normalized log probability.

Therefore, my question is "could we use un-normalized log probability for this function?". My current thought is we use normalized log probability for this function.

To make further clarity, I describe my thought about the difference between un-normalized log probability and normalized log probability.

# logit1 and logit2 are un-normalized log probability for the same distribution.
# Would be logit1 != jax.nn.log_softmax(logit1)
# Would be logit2 != jax.nn.log_softmax(logit2)
jax.nn.log_softmax(logit1) == jax.nn.log_softmax(logit2)

# logit3 is normalized log probability
logit3 == jax.nn.log_softmax(logit3)

Also, I attach the current implementation of jax.random.categorical.

def categorical(key: KeyArray,
                logits: RealArray,
                axis: int = -1,
                shape: Optional[Shape] = None) -> Array:
  """Sample random values from categorical distributions.

  Args:
    key: a PRNG key used as the random key.
    logits: Unnormalized log probabilities of the categorical distribution(s) to sample from,
      so that `softmax(logits, axis)` gives the corresponding probabilities.
    axis: Axis along which logits belong to the same categorical distribution.
    shape: Optional, a tuple of nonnegative integers representing the result shape.
      Must be broadcast-compatible with ``np.delete(logits.shape, axis)``.
      The default (None) produces a result shape equal to ``np.delete(logits.shape, axis)``.

  Returns:
    A random array with int dtype and shape given by ``shape`` if ``shape``
    is not None, or else ``np.delete(logits.shape, axis)``.
  """
  key, _ = _check_prng_key(key)
  check_arraylike("categorical", logits)
  logits_arr = jnp.asarray(logits)

  if axis >= 0:
    axis -= len(logits_arr.shape)

  batch_shape = tuple(np.delete(logits_arr.shape, axis))
  if shape is None:
    shape = batch_shape
  else:
    shape = tuple(shape)
    _check_shape("categorical", shape, batch_shape)

  shape_prefix = shape[:len(shape)-len(batch_shape)]
  logits_shape = list(shape[len(shape) - len(batch_shape):])
  logits_shape.insert(axis % len(logits_arr.shape), logits_arr.shape[axis])
  return jnp.argmax(
      gumbel(key, (*shape_prefix, *logits_shape), logits_arr.dtype) +
      lax.expand_dims(logits_arr, tuple(range(len(shape_prefix)))),
      axis=axis)

[sh0416]

It seems Ok because the difference between the unnormalized log probability is preserved..