Creating A Custom Gradient Function For HMC For Large Datasets

I'm trying to infer the parameters of a Gaussian process using HMC in tensorflow-probability. I have multiple independent data sequences that are generated from the same underlyin

Solution 1:

In case anyone is interested, I was able to solve this by using a custom gradient and stopping the tape recording in the for loop. I needed to import the tape utilities:

from tensorflow.python.eager import tape

then stop recording around the for loop

with tape.stop_recording():
    for i in range(N):
        ...

This stops tracing, I then have to compute the gradient in graph mode and add them manually to stop the OOM error.