def evolve_trotter(psi,
                   H,
                   step_size,
                   num_steps,
                   euclidean=False,
                   callback=None):
    """Evolve an initial wavefunction psi using a trotter decomposition of H.

    If the evolution is euclidean, the wavefunction will be normalized after
    each step.

    Args:
      psi: An `N`-dimensional tensor representing the initial wavefunction.
      H: A list of `N-1` tensors representing nearest-neighbor operators.
      step_size: The trotter step size.
      num_steps: The number of trotter steps to take.
      euclidean: If `True`, evolve in Euclidean (imaginary) time.
      callback: Optional callback function for monitoring the evolution.

    Returns:
      psi_t: The final wavefunction.
      t: The final time.
  """
    num_sites = len(psi.shape)
    layers = trotter_prepare_gates(H, step_size, num_sites, euclidean)
    return _evolve_trotter_gates(psi,
                                 layers,
                                 step_size,
                                 num_steps,
                                 euclidean=euclidean,
                                 callback=callback)
def evolve_trotter_defun(psi,
                         H,
                         step_size,
                         num_steps,
                         euclidean=False,
                         callback=None,
                         batch_size=1):
    """Evolve an initial wavefunction psi using a trotter decomposition of H.

    If the evolution is euclidean, the wavefunction will be normalized after
    each step.

    In this version, `batch_size` steps are "compiled" to a computational graph
    using `defun`, which greatly decreases overhead. 

    Args:
      psi: An `N`-dimensional tensor representing the initial wavefunction.
      H: A list of `N-1` tensors representing nearest-neighbor operators.
      step_size: The trotter step size.
      num_steps: The number of trotter steps to take.
      euclidean: If `True`, evolve in Euclidean (imaginary) time.
      callback: Optional callback function for monitoring the evolution.
      batch_size: The number of steps to unroll in the computational graph.

    Returns:
      psi_t: The final wavefunction.
      t: The final time.
  """
    n_batches, rem = divmod(num_steps, batch_size)

    step_size = tf.cast(step_size, psi.dtype)

    num_sites = len(psi.shape)
    layers = trotter_prepare_gates(H, step_size, num_sites, euclidean)

    t = 0.0
    for i in range(n_batches):
        psi, t_b = _evolve_trotter_gates_defun(psi,
                                               layers,
                                               step_size,
                                               batch_size,
                                               euclidean=euclidean,
                                               callback=None)
        t += t_b
        if callback is not None:
            callback(psi, t, (i + 1) * batch_size - 1)

    if rem > 0:
        psi, t_b = _evolve_trotter_gates_defun(psi,
                                               layers,
                                               step_size,
                                               rem,
                                               euclidean=euclidean,
                                               callback=None)
        t += t_b

    return psi, t
def evolve_trotter(psi, H, step_size, num_steps, euclidean=False, callback=None):
    """Evolve an initial state psi using a trotter decomposition of H.
    If the evolution is euclidean, the wavefunction will be normalized after
    each step.
    """
    num_sites = len(psi.shape)
    layers = trotter_prepare_gates(H, step_size, num_sites, euclidean)
    return _evolve_trotter_gates(
        psi, layers, step_size, num_steps,
        euclidean=euclidean,
        callback=callback)
def evolve_trotter(psi,
                   H,
                   step_size,
                   num_steps,
                   euclidean=False,
                   callback=None):
    """Evolve an initial state psi using a trotter decomposition of H.
    If the evolution is euclidean, the wavefunction will be normalized after
    each step.
    """
    num_sites = len(psi.shape)
    layers = trotter_prepare_gates(H, step_size, num_sites, euclidean)
    return _evolve_trotter_gates(psi,
                                 layers,
                                 step_size,
                                 num_steps,
                                 euclidean=euclidean,
                                 callback=callback)
def evolve_trotter_defun(psi,
                         H,
                         step_size,
                         num_steps,
                         euclidean=False,
                         callback=None,
                         batch_size=1):
    """Evolve an initial state psi using a trotter decomposition of H.
    If the evolution is euclidean, the wavefunction will be normalized after
    each step.

    In this version, `batch_size` steps are "compiled" to a computational graph
    using `defun`, which greatly decreases overhead. 
    """
    n_batches, rem = divmod(num_steps, batch_size)

    step_size = tf.cast(step_size, psi.dtype)

    num_sites = len(psi.shape)
    layers = trotter_prepare_gates(H, step_size, num_sites, euclidean)

    t = 0.0
    for i in range(n_batches):
        psi, t_b = _evolve_trotter_gates_defun(psi,
                                               layers,
                                               step_size,
                                               batch_size,
                                               euclidean=euclidean,
                                               callback=None)
        t += t_b
        if callback is not None:
            callback(psi, t, (i + 1) * batch_size - 1)

    if rem > 0:
        psi, t_b = _evolve_trotter_gates_defun(psi,
                                               layers,
                                               step_size,
                                               rem,
                                               euclidean=euclidean,
                                               callback=None)
        t += t_b

    return psi, t
def evolve_trotter_defun(
    psi, H, step_size, num_steps,
    euclidean=False,
    callback=None,
    batch_size=1):
    """Evolve an initial state psi using a trotter decomposition of H.
    If the evolution is euclidean, the wavefunction will be normalized after
    each step.

    In this version, `batch_size` steps are "compiled" to a computational graph
    using `defun`, which greatly decreases overhead. 
    """
    n_batches, rem = divmod(num_steps, batch_size)

    step_size = tf.cast(step_size, psi.dtype)

    num_sites = len(psi.shape)
    layers = trotter_prepare_gates(H, step_size, num_sites, euclidean)

    t = 0.0
    for i in range(n_batches):
        psi, t_b = _evolve_trotter_gates_defun(
            psi, layers, step_size, batch_size,
            euclidean=euclidean,
            callback=None)
        t += t_b
        if callback is not None:
            callback(psi, t, (i+1) * batch_size - 1)

    if rem > 0:
        psi, t_b = _evolve_trotter_gates_defun(
            psi, layers, step_size, rem, 
            euclidean=euclidean,
            callback=None)
        t += t_b

    return psi, t