Gumbel-Softmax Sample

Problem

A Gumbel-softmax sample perturbs logits with Gumbel noise g_i = -log(-log(u_i)), u_i uniform, then applies a temperature-scaled softmax:

y = softmax((logits + g) / tau)

Use random.Random(seed) and draw the uniforms in order, one per logit. Return the sample vector.

Hint

The Gumbel-max trick makes the argmax of the perturbed logits a categorical sample.

case 1

in: [[0,0],1,42]

expected: [0.8918706788632648,0.10812932113673522]

case 2

in: [[1,2,3],0.5,7]

expected: [0.0026303955863455713,0.0069065892697156425,0.9904630151439388]

case 3

in: [[0,-1],2,123]

expected: [0.5999155108843205,0.40008448911567956]

Starter codePython

def gumbel_softmax_sample(logits, tau, seed):
    # Your code here
    pass

3 test cases · run your code against them in the editor