using a single random key as (unmapped) input to a vmap'ped function.

[hrbigelow]

I wondered if there's any way to indicate semantically that a single unmapped random key should serve as randomness independently to the elements of a mapped dimension?

At the moment, something like jax.random.categorical conveniently allows you to get independent samples across multiple distributions just using a single key. But, if you use such a function from within jax.vmap, is there a way to retain this convenience, or are you stuck now duplicating keys?

# a vector of logits
>>> logits
Array([-0.02830462,  0.46713185,  0.29570296,  0.15354592], dtype=float32)

# each sample is independent even though sampling from same distribution
>>> jax.random.categorical(jax.random.key(42), jnp.broadcast_to(logits, (10, *logits.shape)))
Array([1, 3, 3, 0, 3, 3, 1, 1, 1, 1], dtype=int32)

# each sample identical when calling under jax.vmap.
# unampped `key` argument, but mapped `logits` argument
>>> jax.vmap(jax.random.categorical, in_axes=(None,0))(jax.random.key(42), jnp.broadcast_to(logits, (10, *logits.shape)))
Array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1], dtype=int32)

For reference, jax.random.categorical has this signature

categorical(
    key: 'ArrayLike',
    logits: 'RealArray',
    axis: 'int' = -1,
    shape: 'Shape | None' = None,
    replace: 'bool' = True
) -> 'Array'

[jakevdp]

If you use a single unmapped random key, then the same key will be used within each batch of the computation: this is working by design. If you'd like a different key within each batch, then you can split/fold the key and map over it, e.g. like this:

>>> jax.vmap(jax.random.categorical)(jax.random.split(keys, 10), jnp.broadcast_to(logits, (10, *logits.shape)))
Array([2, 1, 1, 3, 2, 0, 3, 3, 3, 1], dtype=int32)

[hrbigelow]

Yes, I kow that is how it is meant to work. And, on second thought, treating key in some special way by jax.vmap would be a pretty bad idea I think. It's best the way it is.

What I was more broadly wondering about is, are there situations where hand-coding the batched version of a function affords the opportunity for a much more efficient compiled function? Or, does jax.jit + jax.vmap basically always work close to optimally?

[jakevdp]

I don't know of any situations where hand-coding a batched version is much more efficient than just using vmap, but I also wouldn't be surprised if you found cases where that were true. With randomness in particular, there are some samplers (like gamma) that rely on sampling loops that would likely be more efficent drawing samples from a single key vs a batch of keys (because batched while_loop runs until every batch converges).

[hrbigelow]

Ah okay, I'll keep it in mind! I was in the habit of hand-coding a batch dimension in most places, but have become much more comfortable using jax.vmap and mostly using it as a default. Thank you for your advice.