is it possible to use `tfp.optimizer.StochasticGradientLangevinDynamics` with TF2?

Hi,

Is it possible to use tfp.optimizer.StochasticGradientLangevinDynamics with TF2? I was hoping I could try it as a drop in replacement for the classes in tf.optimizers but can't find any way to make it work. Here is my current best attempt, where the OPTIMIZER line can be toggled manually:

import tensorflow as tf
import tensorflow_probability as tfp
import matplotlib.pyplot as plt

A = tf.constant(1.0, dtype=tf.float64)
B = tf.constant(100.0, dtype=tf.float64)
NUM_STEPS = 1000
OPTIMIZER = tf.optimizers.Adam
# OPTIMIZER = tfp.optimizer.StochasticGradientLangevinDynamics


@tf.function
def rosenbrock(x0, x1):
    first = tf.math.squared_difference(A, x0)
    second = B * tf.math.squared_difference(x1, tf.square(x0))
    return first + second


tf.random.set_seed(666)

x0 = tf.Variable(tf.random.normal([], dtype=tf.float64))
x1 = tf.Variable(tf.random.normal([], dtype=tf.float64))

print(f"initial solution: [{x0.numpy():.4f}, {x1.numpy():.4f}]")
print(f"true solution: [{A}, {A*A}]")

optimizer = None


@tf.function
def train(x0, x1):

    global optimizer

    if optimizer is None:
        optimizer = OPTIMIZER(learning_rate=1.0)

    x0s = tf.TensorArray(tf.float64, size=NUM_STEPS)
    x1s = tf.TensorArray(tf.float64, size=NUM_STEPS)

    for step in tf.range(NUM_STEPS):
        optimizer.minimize(lambda: rosenbrock(x0, x1), [x0, x1])

        x0s = x0s.write(step, x0)
        x1s = x1s.write(step, x1)

    return x0s.stack(), x1s.stack()


x0s, x1s = train(x0, x1)

fig, axes = plt.subplots(2)

axes[0].plot(x0s)
axes[0].set_title("x[0]")
axes[0].axhline(A.numpy(), color="red", linestyle="--")

axes[1].plot(x1s)
axes[1].set_title("x[1]")
axes[1].axhline((A * A).numpy(), color="red", linestyle="--")

plt.show()

using SGLD I get:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/jeff/workspace/misc/sgld.py", line 50, in <module>
    x0s, x1s = train(x0, x1)
  File "/home/jeff/.virtualenvs/tf2/lib/python3.6/site-packages/tensorflow_core/python/eager/def_function.py", line 461, in __call__
    return self._stateless_fn(*args, **kwds)
  File "/home/jeff/.virtualenvs/tf2/lib/python3.6/site-packages/tensorflow_core/python/eager/function.py", line 1781, in __call__
    return graph_function._filtered_call(args, kwargs)  # pylint: disable=protected-access
  File "/home/jeff/.virtualenvs/tf2/lib/python3.6/site-packages/tensorflow_core/python/eager/function.py", line 1105, in _filtered_call
    self.captured_inputs)
  File "/home/jeff/.virtualenvs/tf2/lib/python3.6/site-packages/tensorflow_core/python/eager/function.py", line 1183, in _call_flat
    ctx, args, cancellation_manager=cancellation_manager)
  File "/home/jeff/.virtualenvs/tf2/lib/python3.6/site-packages/tensorflow_core/python/eager/function.py", line 475, in call
    ctx=ctx)
  File "/home/jeff/.virtualenvs/tf2/lib/python3.6/site-packages/tensorflow_core/python/eager/execute.py", line 76, in quick_execute
    raise e
  File "/home/jeff/.virtualenvs/tf2/lib/python3.6/site-packages/tensorflow_core/python/eager/execute.py", line 61, in quick_execute
    num_outputs)
TypeError: An op outside of the function building code is being passed
a "Graph" tensor. It is possible to have Graph tensors
leak out of the function building context by including a
tf.init_scope in your function building code.
For example, the following function will fail:
  @tf.function
  def has_init_scope():
    my_constant = tf.constant(1.)
    with tf.init_scope():
      added = my_constant * 2
The graph tensor has name: StochasticGradientLangevinDynamics/learning_rate:0

But with Adam it works as expected.

Are there any workarounds? I would be happy to invest some time into adding a fix but would need a little initial guidance.

Thanks!

>>> tfp.__version__
'0.8.0-dev20190814'
>>> tf.__version__
'2.0.0-rc0'

Reproduce the provided Rosenbrock example with TensorFlow Probability 0.8.0-dev20190814 and TensorFlow 2.0.0-rc0, comparing StochasticGradientLangevinDynamics with tf.optimizers.Adam inside tf.function. Start at tfp.optimizer.StochasticGradientLangevinDynamics and trace the learning_rate tensor handling; done means the optimizer works in the shown TF2 training loop without the graph-tensor error.