Decouple sampling from pseudo_kl_divergence_loss. (#34)

Co-authored-by: Vladimir Vargas Calderón <vvargasc@dwavesys.com>

Author: VolodyaCO

dwave/plugins/torch/models/losses/kl_divergence.py

@@ -12,10 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
 
 import torch
-from dimod import Sampler
 
 from dwave.plugins.torch.models.boltzmann_machine import GraphRestrictedBoltzmannMachine
 
@@ -25,12 +23,8 @@
 def pseudo_kl_divergence_loss(
     spins: torch.Tensor,
     logits: torch.Tensor,
+    samples: torch.Tensor,
     boltzmann_machine: GraphRestrictedBoltzmannMachine,
-    sampler: Sampler,
-    sample_kwargs: dict,
-    prefactor: Optional[float] = None,
-    linear_range: Optional[tuple[float, float]] = None,
-    quadratic_range: Optional[tuple[float, float]] = None,
 ):
     """A pseudo Kullback-Leibler divergence loss function for a discrete autoencoder with a
     Boltzmann machine prior.
@@ -47,32 +41,11 @@ def pseudo_kl_divergence_loss(
             logits are the raw output of the encoder.
         boltzmann_machine (GraphRestrictedBoltzmannMachine): An instance of a Boltzmann
             machine.
-        sampler (Sampler): A sampler used for generating samples.
-        sample_kwargs (dict): Additional keyword arguments for the ``sampler.sample``
-            method.
-        prefactor (float, optional): A scaling applied to the Hamiltonian weights
-            (linear and quadratic weights). When None, no scaling is applied. Defaults
-            to None.
-        linear_range (tuple[float, float], optional): Linear weights are clipped to
-            ``linear_range`` prior to sampling. This clipping occurs after the
-            ``prefactor`` scaling has been applied. When None, no clipping is applied.
-            Defaults to None.
-        quadratic_range (tuple[float, float], optional): Quadratic weights are clipped
-            to ``quadratic_range`` prior to sampling. This clipping occurs after the
-            ``prefactor`` scaling has been applied. When None, no clipping is applied.
-            Defaults to None.
+        samples (torch.Tensor): A tensor of samples from the Boltzmann machine.
 
     Returns:
         torch.Tensor: The computed pseudo KL divergence loss.
     """
-    samples = boltzmann_machine.sample(
-        sampler=sampler,
-        device=spins.device,
-        prefactor=prefactor if prefactor is not None else 1.0,
-        linear_range=linear_range,
-        quadratic_range=quadratic_range,
-        sample_params=sample_kwargs,
-    )
     probabilities = torch.sigmoid(logits)
     entropy = torch.nn.functional.binary_cross_entropy_with_logits(logits, probabilities)
     cross_entropy = boltzmann_machine.quasi_objective(spins, samples)

releasenotes/notes/kld-requires-boltzmann-samples-67ae24105364b205.yaml

@@ -0,0 +1,7 @@
+---
+upgrade:
+  - |
+    ``dwave.plugins.torch.models.losses.kl_divergence.pseudo_kl_divergence_loss``
+    no longer uses the Graph-Restricted Boltzmann Machine to generate Boltzmann
+    samples internally. Instead, the samples must be provided as an argument
+    to the function. This is a breaking change.

tests/test_dvae_winci2020.py

@@ -129,12 +129,17 @@ def test_train(self, n_latent_dims):
 
             discretes = discretes.reshape(discretes.shape[0], -1)
             latents = latents.reshape(latents.shape[0], -1)
+            samples = self.boltzmann_machine.sample(
+                self.sampler_sa,
+                as_tensor=True,
+                prefactor=1.0,
+                sample_params=dict(num_sweeps=10, seed=1234, num_reads=100),
+            )
             kl_loss = pseudo_kl_divergence_loss(
                 discretes,
                 latents,
+                samples,
                 self.boltzmann_machine,
-                self.sampler_sa,
-                dict(num_sweeps=10, seed=1234, num_reads=100),
             )
             loss = loss + 1e-1 * kl_loss
             optimiser.zero_grad()