import tensorflow as tf


  for _ in range(n_iters):
    _, samples = step_fn(dynamics, optimizer, samples)


def fit(dynamics,
        samples,
        optimizer,
        step_fn=step,
        n_iters=5000,
        verbose=True,
        logdir=None):
  """Fit L2HMC sampler with given log-likelihood function."""

  if logdir:
    summary_writer = tf.contrib.summary.create_file_writer(logdir)

  for i in range(n_iters):
    loss, samples = step_fn(dynamics, optimizer, samples)
    if verbose:
      print("Iteration %d: loss %.4f" % (i, loss))

    if logdir:
      with summary_writer.as_default():
        with tf.contrib.summary.always_record_summaries():
          tf.contrib.summary.scalar("loss", loss)


class L2hmcTest(tf.test.TestCase):
  """Unit tests for l2hmc in both eager and graph mode."""
