Langevin Step

Problem

Take one Langevin dynamics step on a scalar parameter.

After random.seed(seed), return x - lr * grad + noise_std * sqrt(lr) * random.gauss(0.0, 1.0). The noise is scaled by the square root of the step size.

Hint

Reseed so the injected Gaussian draw is reproducible.

Test cases

case 1

in: [1,0.5,0.1,1,0]

expected: 1.2477965586458089

case 2

in: [0,-1,0.25,0.5,42]

expected: 0.21397741760551792

case 3

in: [2,0,0.5,2,7]

expected: 1.6381306257334638

Starter codePython

def langevin_step(x, grad, lr, noise_std, seed):
    # Your code here
    pass

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