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


    hparams = get_default_hparams()
    tf.enable_resource_variables()
    for sample_size in [10, 25, 50, 100, 200]:
      hparams.n_samples = sample_size
      tf.reset_default_graph()
      with tf.Graph().as_default():
        energy_fn, _, _ = l2hmc.get_scg_energy_fn()
        x = tf.random_normal([hparams.n_samples, hparams.x_dim],
                             dtype=tf.float32)
        dynamics = l2hmc.Dynamics(
            x_dim=hparams.x_dim,
            minus_loglikelihood_fn=energy_fn,
            n_steps=hparams.n_steps,
            eps=hparams.eps)
        loss, _, _ = l2hmc.compute_loss(dynamics, x)

        optimizer = tf.train.AdamOptimizer(learning_rate=hparams.learning_rate)
        train_op, loss, _ = graph_step(dynamics, optimizer, x)

        # Single thread; fairer comparison against eager
        session_conf = tf.ConfigProto(inter_op_parallelism_threads=1)

        with tf.Session(config=session_conf) as sess:
          sess.run(tf.global_variables_initializer())

          # Warmup to reduce initialization effect when timing
          for _ in range(hparams.n_warmup_iters):
            _, _ = sess.run([train_op, loss])

