QHBMs as second argument in VQT and QMHL

qhbmlib/ebm.py

@@ -547,6 +547,14 @@ def has_operator(self):
   def is_analytic(self):
     return self._is_analytic
 
+  @property
+  def trainable_variables(self):
+    return self._energy_function.trainable_variables
+
+  @trainable_variables.setter
+  def trainable_variables(self, value):
+    self._energy_function.trainable_variables = value
+
   def copy(self):
     if self._energy_sampler is not None:
       energy_sampler = self._energy_sampler.copy()
@@ -558,7 +566,7 @@ def copy(self):
         energy_function,
         energy_sampler,
         is_analytic=self.is_analytic,
-        name=self.name)
+        name=f'{self.name}_copy')
 
   def energy(self, bitstrings):
     return self._energy_function.energy(bitstrings)
@@ -588,7 +596,7 @@ def energies(self):
 
   def probabilities(self):
     if self.is_analytic:
-      return tf.exp(-self.ebm.energies()) / tf.exp(
+      return tf.exp(-self.energies()) / tf.exp(
           self.log_partition_function())
     raise NotImplementedError()
 
@@ -645,7 +653,7 @@ def is_analytic(self):
 
   def copy(self):
     bernoulli = Bernoulli(
-        self.num_bits, is_analytic=self.is_analytic, name=self.name)
+        self.num_bits, is_analytic=self.is_analytic, name=f'{self.name}_copy')
     bernoulli.kernel.assign(self.kernel)
     return bernoulli
 

qhbmlib/qhbm.py

@@ -72,22 +72,69 @@ def is_analytic(self):
     return self.ebm.is_analytic and self.qnn.is_analytic
 
   def copy(self):
-    return QHBM(self.ebm.copy(), self.qnn.copy(), name=self.name)
+    return QHBM(self.ebm.copy(), self.qnn.copy(), name=f'{self.name}_copy')
 
-  def circuits(self, num_samples):
-    bitstrings, counts = self.ebm.sample(num_samples)
-    circuits = self.qnn.circuits(bitstrings)
-    return circuits, counts
+  def circuits(self, num_samples, unique=True, resolve=True):
+    if unique:
+      bitstrings, counts = self.ebm.sample(num_samples, unique=unique)
+      circuits = self.qnn.circuits(bitstrings, resolve=resolve)
+      return circuits, counts
+    bitstrings = self.ebm.sample(num_samples, unique=unique)
+    circuits = self.qnn.circuits(bitstrings, resolve=resolve)
+    return circuits
 
   def sample(self, num_samples, mask=True, reduce=True, unique=True):
     bitstrings, counts = self.ebm.sample(num_samples)
     return self.qnn.sample(
         bitstrings, counts, mask=mask, reduce=reduce, unique=unique)
 
-  def expectation(self, operators, num_samples, reduce=True):
+  def expectation(self, operators, num_samples, mask=True, reduce=True):
+    """TODO: add gradient function"""
+    if isinstance(operators, QHBM):
+      circuits, counts = self.circuits(num_samples, resolve=False)
+      return operators.operator_expectation((circuits, counts),
+                                            symbol_names=self.qnn.symbols,
+                                            symbol_values=self.qnn.values,
+                                            mask=mask,
+                                            reduce=reduce)
     bitstrings, counts = self.ebm.sample(num_samples)
     return self.qnn.expectation(bitstrings, counts, operators, reduce=reduce)
 
+  def operator_expectation(self,
+                           density_operator,
+                           num_samples=None,
+                           symbol_names=None,
+                           symbol_values=None,
+                           reduce=True,
+                           mask=True):
+    """TODO: add gradient function"""
+    if isinstance(density_operator, tuple):
+      circuits, counts = density_operator
+    elif isinstance(density_operator, QHBM):
+      circuits, counts = density_operator.circuits(num_samples, resolve=False)
+      symbol_names = density_operator.qnn.symbols
+      symbol_values = density_operator.qnn.values
+    else:
+      raise TypeError()
+
+    if self.ebm.has_operator:
+      expectation_shards = self.qnn.pulled_back_expectation(
+          circuits,
+          counts,
+          self.operator_shards,
+          symbol_names=symbol_names,
+          symbol_values=symbol_values,
+          reduce=reduce)
+      return self.ebm.operator_expectation(expectation_shards)
+    bitstrings, counts = self.qnn.pulled_back_sample(
+        circuits, counts, mask=mask)
+    energies = self.ebm.energy(bitstrings)
+    if reduce:
+      probs = tf.cast(counts, tf.float32) / tf.cast(
+          tf.reduce_sum(counts), tf.float32)
+      return tf.reduce_sum(probs * energies)
+    return energies
+
   def probabilities(self):
     return self.ebm.probabilities()
 
@@ -111,10 +158,12 @@ def density_matrix(self):
 
   def fidelity(self, sigma: tf.Tensor):
     """TODO: convert to tf.keras.metric.Metric
+
     Calculate the fidelity between a QHBM and a density matrix.
         Args:
           sigma: 2-D `tf.Tensor` of dtype `complex64` representing the right
             density matrix in the fidelity calculation.
+
         Returns:
           A scalar `tf.Tensor` which is the fidelity between the density matrix
             represented by this QHBM and `sigma`.

qhbmlib/qmhl.py

@@ -15,18 +15,19 @@
 """Impementations of the QMHL loss and its derivatives."""
 
 import tensorflow as tf
+from qhbmlib import qhbm
 
 
-def qmhl(qhbm_model, target_circuits, target_counts):
+def qmhl(qhbm_model, density_operator, num_samples=1000):
   """Calculate the QMHL loss of the qhbm model against the target.
 
   This loss is differentiable with respect to the trainable variables of the
     model.
 
   Args:
     qhbm_model: Parameterized model density operator.
-    target_circuits: 1-D tensor of strings which are serialized circuits.
-      These circuits represent samples from the data density matrix.
+    target_circuits: 1-D tensor of strings which are serialized circuits. These
+      circuits represent samples from the data density matrix.
     target_counts: 1-D tensor of integers which are the number of samples to
       draw from the data density matrix: `target_counts[i]` is the number of
         samples to draw from `target_circuits[i]`.
@@ -36,73 +37,77 @@ def qmhl(qhbm_model, target_circuits, target_counts):
   """
 
   @tf.custom_gradient
-  def loss(trainable_variables):
-    # log_partition estimate
+  def function(trainable_variables):
+    # pulled back expectation of energy operator
+    if isinstance(density_operator, tuple):
+      circuits, counts = density_operator
+    elif isinstance(density_operator, qhbm.QHBM):
+      circuits, counts = density_operator.circuits(num_samples)
+    else:
+      raise TypeError()
 
+    if qhbm_model.ebm.has_operator:
+      expectation_shards = qhbm_model.qnn.pulled_back_expectation(
+          circuits, counts, qhbm_model.operator_shards)
+      expectation = qhbm_model.ebm.operator_expectation(expectation_shards)
+    else:
+      qnn_bitstrings, qnn_counts = qhbm_model.qnn.pulled_back_sample(
+          circuits, counts)
+      energies = qhbm_model.ebm.energy(qnn_bitstrings)
+      qnn_probs = tf.cast(qnn_counts, tf.float32) / tf.cast(
+          tf.reduce_sum(qnn_counts), tf.float32)
+      expectation = tf.reduce_sum(qnn_probs * energies)
+
+    # log_partition estimate
     if qhbm_model.ebm.is_analytic:
       log_partition_function = qhbm_model.log_partition_function()
     else:
-      bitstrings, _ = qhbm_model.ebm.sample(tf.reduce_sum(target_counts))
+      bitstrings, _ = qhbm_model.ebm.sample(tf.reduce_sum(counts))
       energies = qhbm_model.ebm.energy(bitstrings)
-      log_partition_function = tf.math.reduce_logsumexp(-1 * energies)
+      log_partition_function = tf.math.reduce_logsumexp(-energies)
 
-    # pulled back expectation of energy operator
-    qnn_bitstrings, qnn_counts = qhbm_model.qnn.pulled_back_sample(
-        target_circuits, target_counts)
-    qnn_probs = tf.cast(qnn_counts, tf.float32) / tf.cast(
-        tf.reduce_sum(qnn_counts), tf.float32)
-    energies = qhbm_model.ebm.energy(qnn_bitstrings)
-    avg_energy = tf.reduce_sum(qnn_probs * energies)
-
-    def grad(grad_y, variables=None):
+    def gradient(grad_y, variables=None):
       """Gradients are computed using estimators from the QHBM paper."""
-      # Thetas derivative.
-      ebm_bitstrings, ebm_counts = qhbm_model.ebm.sample(
-          tf.reduce_sum(target_counts))
-      ebm_probs = tf.cast(ebm_counts, tf.float32) / tf.cast(
-          tf.reduce_sum(ebm_counts), tf.float32)
-      with tf.GradientTape() as tape:
+      with tf.GradientTape(persistent=True) as tape:
         tape.watch(qhbm_model.ebm.trainable_variables)
-        qnn_energies = qhbm_model.ebm.energy(qnn_bitstrings)
-      # jacobian is a list over thetas, with ith entry a tensor of shape
-      # [tf.shape(qnn_energies)[0], tf.shape(thetas[i])[0]]
-      qnn_energy_jac = tape.jacobian(qnn_energies,
-                                     qhbm_model.ebm.trainable_variables)
+        if qhbm_model.ebm.has_operator:
+          tape.watch(qhbm_model.qnn.trainable_variables)
+          expectation_shards = qhbm_model.qnn.pulled_back_expectation(
+              circuits, counts, qhbm_model.operator_shards)
+          expectation = qhbm_model.ebm.operator_expectation(expectation_shards)
+        else:
+          energies = qhbm_model.ebm.energy(qnn_bitstrings)
+          expectation = tf.reduce_sum(qnn_probs * energies)
+      qnn_energy_grad = tape.gradient(expectation,
+                                      qhbm_model.ebm.trainable_variables)
+      if qhbm_model.ebm.has_operator:
+        grad_qnn = tape.gradient(expectation,
+                                 qhbm_model.qnn.trainable_variables)
+        grad_qnn = [grad_y * grad for grad in grad_qnn]
+      else:
+        raise NotImplementedError(
+            "Derivative when EBM has no operator is not yet supported.")
 
+      ebm_bitstrings, ebm_counts = qhbm_model.ebm.sample(tf.reduce_sum(counts))
+      ebm_probs = tf.cast(ebm_counts, tf.float32) / tf.cast(
+          tf.reduce_sum(ebm_counts), tf.float32)
       with tf.GradientTape() as tape:
         tape.watch(qhbm_model.ebm.trainable_variables)
-        ebm_energies = qhbm_model.ebm.energy(ebm_bitstrings)
-      ebm_energy_jac = tape.jacobian(ebm_energies,
-                                     qhbm_model.ebm.trainable_variables)
+        energies = qhbm_model.ebm.energy(ebm_bitstrings)
+        expectation = tf.reduce_sum(ebm_probs * energies)
+      ebm_energy_grad = tape.gradient(expectation,
+                                      qhbm_model.ebm.trainable_variables)
 
-      # contract over bitstring weights
       grad_ebm = [
-          grad_y *
-          (tf.reduce_sum(
-              tf.transpose(qnn_probs * tf.transpose(qnn_energy_grad)), 0) -
-           tf.reduce_sum(
-               tf.transpose(ebm_probs * tf.transpose(ebm_energy_grad)), 0))
-          for qnn_energy_grad, ebm_energy_grad in zip(qnn_energy_jac,
-                                                      ebm_energy_jac)
+          grad_y * (qnn_grad - ebm_grad)
+          for qnn_grad, ebm_grad in zip(qnn_energy_grad, ebm_energy_grad)
       ]
 
-      # Phis derivative.
-      if qhbm_model.ebm.has_operator:
-        with tf.GradientTape() as tape:
-          tape.watch(qhbm_model.qnn.trainable_variables)
-          energy_shards = qhbm_model.qnn.pulled_back_expectation(
-              target_circuits, target_counts, qhbm_model.operator_shards)
-          energy = qhbm_model.ebm.operator_expectation(energy_shards)
-        grad_qnn = tape.gradient(energy, qhbm_model.qnn.trainable_variables)
-        grad_qnn = [grad_y * g for g in grad_qnn]
-      else:
-        raise NotImplementedError(
-            "Derivative when EBM has no operator is not yet supported.")
       grad_qhbm = grad_ebm + grad_qnn
-      if variables is None:
-        return grad_qhbm
-      return grad_qhbm, [tf.zeros_like(g) for g in grad_qhbm]
+      if variables:
+        return grad_qhbm, [tf.zeros_like(var) for var in variables]
+      return grad_qhbm
 
-    return avg_energy + log_partition_function, grad
+    return expectation + log_partition_function, gradient
 
-  return loss(qhbm_model.trainable_variables)
+  return function(qhbm_model.trainable_variables)