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



# To be defunnable, the function cannot return an Operation, so the above
# function is used for defun or eager, and this function is used in graph to be
# able to run the gradient updates.
def graph_step(dynamics, optimizer, samples):
  loss, grads, samples, _ = l2hmc.loss_and_grads(
      dynamics, samples, loss_fn=l2hmc.compute_loss)
  train_op = optimizer.apply_gradients(zip(grads, dynamics.variables))

  return train_op, loss, samples

