Straight-Through Gumbel

Problem

The straight-through Gumbel estimator uses a hard one-hot sample in the forward pass while backpropagating through the soft sample.

Generate the same seeded Gumbel-softmax sample, then take the one-hot vector of its argmax (first maximum wins). Return [hard_onehot, soft_sample].

Hint

In the backward pass the hard sample is treated as if it were the soft one.

case 1

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

expected: [[1,0],[0.8918706788632648,0.10812932113673522]]

case 2

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

expected: [[0,0,1],[0.0026303955863455713,0.0069065892697156425,0.9904630151439388]]

case 3

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

expected: [[1,0],[0.5999155108843205,0.40008448911567956]]

Starter codePython

def straight_through_gumbel(logits, tau, seed):
    # Return [hard_onehot, soft_sample]
    # Your code here
    pass

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