from tensorflow.contrib.eager.python.examples.l2hmc import l2hmc


  def benchmark_eager_defun(self):
    self._benchmark_eager(defun=True)

  def _benchmark_eager(self, defun=False):
    """Benchmark Eager performance."""

    hparams = get_default_hparams()
    for sample_size in [10, 25, 50, 100, 200]:
      hparams.n_samples = sample_size
      energy_fn, _, _ = l2hmc.get_scg_energy_fn()
      dynamics = l2hmc.Dynamics(
          x_dim=hparams.x_dim,
          minus_loglikelihood_fn=energy_fn,
          n_steps=hparams.n_steps,
          eps=hparams.eps)
      optimizer = tf.train.AdamOptimizer(learning_rate=hparams.learning_rate)
      step_fn = tfe.defun(step) if defun else step

      # Warmup to reduce initialization effect when timing
