Add conditional sampling by anahitamansouri · Pull Request #68 · dwavesystems/dwave-pytorch-plugin

anahitamansouri · 2026-03-09T17:18:03Z

This PR adds:

Conditional sampling feature for block spin sampling.
BipartiteSampler for sampling bipartite GRBMs.
An example of using the BipartiteSampler.
Tests for the new functionalities.

kevinchern

Excellent!!!! Did a quick first pass with minor requests. The implementation is clean and efficient, documentation is well written, and tests are thorough.
Missing implementation for DimodSampler but let's add that in a separate PR.

kevinchern

Added a couple minor requests

kevinchern · 2026-03-14T01:27:50Z

+            torch.Tensor: A tensor of shape (num_chains, n_nodes) of +/-1 values sampled from the model.
+        """
+        if x is not None:
+            mask = self._validate_input_and_generate_mask(x)


Suggested change

mask = self._validate_input_and_generate_mask(x)

self._validate_input(x)

mask = ~torch.isnan(x)

kevinchern · 2026-03-14T01:28:11Z

+            h = self._grbm.hidden_idx
+            self._x[:, h] = torch.where(mask[:, h], x[:, h], self._x[:, h])
+
+    def _validate_input_and_generate_mask(self, x: torch.Tensor) -> torch.Tensor:


Suggested change

def _validate_input_and_generate_mask(self, x: torch.Tensor) -> torch.Tensor:

def _validate_input(self, x: torch.Tensor) -> None:

kevinchern · 2026-03-14T01:29:39Z

+            self._x[:, h] = torch.where(mask[:, h], x[:, h], self._x[:, h])
+
+    def _validate_input_and_generate_mask(self, x: torch.Tensor) -> torch.Tensor:
+        """Validate conditional sampling input and construct a boolean mask.


Suggested change

"""Validate conditional sampling input and construct a boolean mask.

"""Validate conditional sampling input.

kevinchern · 2026-03-14T01:30:49Z

+
+        Returns:
+            torch.Tensor: Boolean mask of shape ``(num_chains, n_nodes)`` where
+            ``True`` indicates clamped variables (observed in ``x``) and
+            ``False`` indicates variables that should be sampled (``NaN`` in x).


Suggested change

Returns:

torch.Tensor: Boolean mask of shape ``(num_chains, n_nodes)`` where

``True`` indicates clamped variables (observed in ``x``) and

``False`` indicates variables that should be sampled (``NaN`` in x).

kevinchern · 2026-03-14T01:31:00Z

+                "The input must be unclamped for visible or hidden but not both."
+            )
+
+        return mask


Suggested change

return mask

kevinchern · 2026-03-14T01:32:25Z

+
+        Args:
+            x (torch.Tensor): A tensor of shape (``num_chains``, ``dim``) or (``num_chains``, ``n_nodes``)
+                interpreted as a batch of partially-observed spins. Entries marked with ``torch.nan`` will


Suggested change

interpreted as a batch of partially-observed spins. Entries marked with ``torch.nan`` will

interpreted as a batch of partially observed spins. Entries marked with ``torch.nan`` will

kevinchern · 2026-03-14T01:34:55Z

+            if mask is not None:
+                self._x[:, block] = torch.where(mask[:, block], x[:, block], self._x[:, block])
+
+    def _validate_input_and_generate_mask(self, x: torch.Tensor) -> torch.Tensor:


Same suggestion here as in bipartite sampler (docstring, type hints, returns, and defining mask outside)

thisac

I recall we talked about this, but it seems like BipartiteGibbsSampler and BlockSampler could share a lot of methods and do with some deduplication. If they're not general enough to fit into TorchSampler, there should either be a hierarchy between them or another common class that they inherit from, or, especially if you foresee some of these methods being used in other samplers, you could create one (or several) mixin classes.

kevinchern

Almost there!

kevinchern · 2026-03-23T18:55:53Z

+        grbm = GRBM(nodes, edges, hidden_nodes=["h1", "h2"])
+
+        def crayon(n):
+            return 0 if n in ["v1", "v2"] else 1


Suggested change

return 0 if n in ["v1", "v2"] else 1

return n in ["v1", "v2"]

kevinchern · 2026-03-23T19:02:15Z

+            if mask is not None:
+                self._x[:, block] = torch.where(mask[:, block], x[:, block], self._x[:, block])
+
+    def _validate_input_and_generate_mask(self, x: torch.Tensor) -> torch.Tensor:


Bumping this suggestion to separate validation and mask generation

kevinchern · 2026-03-23T19:06:43Z

+        self.assertEqual(mask.shape, x_valid.shape)
+
+        # Chain 0: visible unclamped
+        self.assertTrue(mask[0, 2:].all()) # First chain: hidden spins are clamped
+
+        # Chain 1: hidden unclamped
+        self.assertTrue(mask[1, :2].all()) # Second chain: visible spins are clamped


IF we keep the signature of validate_and_generate..., THEN an we combine these tests into one where the mask is hard-coded like expected_mask = torch.tensor([[False, ...], [...]])?

e.g., torch.testing.assertEqual(mask, expected_mask) or self.assertListEqual(mask.tolist(), expected_mask.tolist())

kevinchern · 2026-03-23T19:46:45Z

+        # Gibbs update for hidden block (block=1)
+        with self.subTest("hidden block Gibbs update"):
+            sampler._gibbs_update(0.0, hidden_block, ones*zero_field)
+            torch.testing.assert_close(torch.tensor(0.0), sampler._x.mean(), atol=1e-2, rtol=1e-2)


Why does this one have a looser tolerance 1e-2 than the previous 1e-3?

Yeah, I noticed that there are fewer random variables in this test compared to the above one, so the estimate has higher variance, and I needed a looser tolerance (1e-2). I could avoid this by setting sampler._x.data[:] = 1.0 just like the earlier example. I can update the test if you think so.

I might be missing something---don't both tests use sampler._x.mean() so the sample size should be the same(?)

kevinchern · 2026-03-23T19:53:38Z

+    def test_sample_conditional(self):
+        nodes = ["v1", "v2", "h1", "h2"]
+        edges = [["v1", "h1"], ["v1", "h2"], ["v2", "h1"], ["v2", "h2"]]
+        grbm = GRBM(nodes, edges, hidden_nodes=["h1", "h2"])


consider setting the linear fields to be very large so the result is ~= deterministic.
Then, in the test cases, hard-code the expected results per conditional sampling step.

e.g.,
grbm.linear.data[:] = 99999999999999
and sampler._x.data[:] = 1
in one conditional step, everything but the clamped-states should become -1.

kevinchern · 2026-03-30T22:57:56Z

+    Add conditional sampling functionality for the ``BlockSampler``.
+  - |
+    Add ``.clone()`` to the return of ``BlockSampler.sample`` to prevent
+    unintended in-place modification of the sampler’s internal state due to


Suggested change

unintended in-place modification of the sampler’s internal state due to

unintended in-place modification of the sampler's internal state due to

thisac

Just a couple of minor comments.

thisac · 2026-04-13T22:29:24Z

+    from dwave.plugins.torch.models.boltzmann_machine import (
+        GraphRestrictedBoltzmannMachine as GRBM,
+    )
+    from torch._prims_common import DeviceLikeType


Better avoid importing from "private" module and either have the type alias declared here directly (DeviceLikeType: TypeAlias = str | torch.device | int) or just use either the combined type or *args, **kwargs in the signature as PyTorch does for the to() method.

have the type alias declared here directly

@thisac what does this mean 🤔? Is this something you declare at import?

You could just add

DeviceLikeType: TypeAlias = str | torch.device | int

to the top of the file (or any file, like utils.py or base.py, and import it from there).

thisac · 2026-04-13T22:35:31Z

+        """
+        Computes the effective field for all vertices in ``block``.


Suggested change

"""

Computes the effective field for all vertices in ``block``.

"""Computes the effective field for all vertices in ``block``.

thisac · 2026-04-13T23:12:37Z

+            mask (torch.Tensor, optional): Boolean tensor of shape 
+                ``(num_chains, n_nodes)`` indicating which variables are clamped. 
+                Entries set to ``True`` will keep their values during sampling.


Could this be named something clarifying, like clamp_mask, instead of generically mask?

thisac · 2026-04-13T23:15:01Z

+            ValueError: If ``x`` does not match the sampler state shape
+            ``(num_chains, n_nodes)``, contains values other than ``±1``
+            or ``NaN``, or if both visible and hidden variables are
+            simultaneously unclamped within the same chain.


Raises should be indented similarly to Args (not Returns).

Suggested change

ValueError: If ``x`` does not match the sampler state shape

``(num_chains, n_nodes)``, contains values other than ``±1``

or ``NaN``, or if both visible and hidden variables are

simultaneously unclamped within the same chain.

ValueError: If ``x`` does not match the sampler state shape

``(num_chains, n_nodes)``, contains values other than ``±1``

or ``NaN``, or if both visible and hidden variables are

simultaneously unclamped within the same chain.

thisac · 2026-04-13T23:15:44Z

+            mask = None
+        for beta in self._schedule:
+            self._step(beta, mask=mask, x=x)
+        return self._x.clone()


Final empty line missing.

Suggested change

return self._x.clone()

return self._x.clone()

I never bothered to google why this convention is adopted and simply trusted my autoformatter 😆.
Just googled it and thought i'd share a few here (omitted a couple):

@kevinchern what autoformatter are you using?

@anahitamansouri I'd recommend using Black with a line-length set to 100 (black -l 100). Just make sure to only touch files and lines that you've added.

Yeah, I used that on our code in a PR and noticed it was changing some lines that weren't my changes. So, I thought people are not using Black here. So, I did not use that anymore :) I'll try it again with 100.

It's the recommended tool, but we don't enforce it, so there's a fair bit of code that isn't formatted accordingly.

sry misesd the notification @anahitamansouri, I use autopep8 with some customization following dwave's contributor guidelinees

thisac · 2026-04-13T23:17:14Z

+        else:
+            mask = None
+        for beta in self._schedule:
+            self._step(beta, mask=mask, x=x)


IMO makes it much more legible to separate the if-else from for and return with empty lines.

Suggested change

else:

mask = None

for beta in self._schedule:

self._step(beta, mask=mask, x=x)

else:

mask = None

for beta in self._schedule:

self._step(beta, mask=mask, x=x)

thisac · 2026-04-13T23:19:19Z


        Args:
            beta (torch.Tensor): Inverse temperature to sample at.
+            mask (torch.Tensor, optional): Boolean tensor of shape 


Same comment here re renaming.

thisac · 2026-04-13T23:22:31Z

+            mask = None
        for beta in self._schedule:
-            self._step(beta)
-        return self._x
+            self._step(beta, mask, x)
+        return self._x.clone()


Same comment here also about empty lines between if-else and for.

thisac · 2026-04-13T23:30:21Z

+  - |
+    Add ``.clone()`` to the return of ``BlockSampler.sample`` to prevent
+    unintended in-place modification of the sampler's internal state due to
+    returning a reference to the underlying tensor.


Less a feature and more a fix or upgrade, no?

Actually I thought it's a mix of both. I thought BipartiteSampler is a feature and conditional sampling is an upgrade. How would you characterize a feature? :)

So do you suggest to move both under upgrade?

Ah, sorry if I was unclear. I only meant moving the last bullet. The other two are fine.

Ah sorry. This makes sense :)

thisac · 2026-04-13T23:34:38Z

+    def test_prepare_initial_states(self):
+        nodes = ["v1", "v2", "h1", "h2"]
+        edges = [["v1", "h1"], ["v1", "h2"], ["v2", "h1"], ["v2", "h2"]]
+        grbm = GRBM(nodes, edges, hidden_nodes=["h1", "h2"])
+
+        sampler = BipartiteGibbsSampler(grbm, num_chains=2, schedule=[1.0],)
+        # Invalid spins
+        with self.subTest("Non-spin initial states."):
+            self.assertRaisesRegex(ValueError, "contain nonspin values", sampler._prepare_initial_states,
+                            initial_states=torch.tensor([[0, 1, -1, 1]]), num_chains=1)
+
+        # Incorrect shape
+        with self.subTest("Testing initial states with incorrect shape."):
+            self.assertRaisesRegex(ValueError, "Initial states should be of shape", sampler._prepare_initial_states,
+                              num_chains=2, initial_states=torch.tensor([[-1, 1, 1, 1, -1]]))


Only tests exceptions raised. Perhaps rename this test_prepare_initial_states_exceptions and have another test_prepare_initial_states with a test for valid arguments.

Add conditional sampling

ac9c328

anahitamansouri requested a review from kevinchern March 9, 2026 17:18

anahitamansouri self-assigned this Mar 9, 2026

anahitamansouri added the enhancement New feature or request label Mar 9, 2026

Add release note for conditional sampling

ef5b1dc

anahitamansouri marked this pull request as ready for review March 9, 2026 21:18

kevinchern requested changes Mar 9, 2026

View reviewed changes

Addressing the review comments

b4274e3

kevinchern requested changes Mar 17, 2026

View reviewed changes

thisac reviewed Mar 18, 2026

View reviewed changes

kevinchern requested changes Mar 23, 2026

View reviewed changes

Improving samplers, tests and release note

6a52d94

kevinchern self-requested a review March 26, 2026 18:40

kevinchern approved these changes Mar 30, 2026

View reviewed changes

Fixing the apostrophe

fc55216

anahitamansouri requested a review from thisac April 2, 2026 22:46

thisac reviewed Apr 13, 2026

View reviewed changes

improve formatting and test_prepare_initial_states

4ec4fc5

	mask = self._validate_input_and_generate_mask(x)
	self._validate_input(x)
	mask = ~torch.isnan(x)

	def _validate_input_and_generate_mask(self, x: torch.Tensor) -> torch.Tensor:
	def _validate_input(self, x: torch.Tensor) -> None:

	"""Validate conditional sampling input and construct a boolean mask.
	"""Validate conditional sampling input.

	interpreted as a batch of partially-observed spins. Entries marked with ``torch.nan`` will
	interpreted as a batch of partially observed spins. Entries marked with ``torch.nan`` will

	return 0 if n in ["v1", "v2"] else 1
	return n in ["v1", "v2"]

	unintended in-place modification of the sampler’s internal state due to
	unintended in-place modification of the sampler's internal state due to

		"""
		Computes the effective field for all vertices in ``block``.

	"""
	Computes the effective field for all vertices in ``block``.
	"""Computes the effective field for all vertices in ``block``.

Conversation

anahitamansouri commented Mar 9, 2026

Uh oh!

kevinchern left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

kevinchern left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

thisac left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

kevinchern left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

anahitamansouri Mar 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

thisac left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kevinchern Apr 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

kevinchern left a comment •

edited

Loading

anahitamansouri Mar 24, 2026 •

edited

Loading

kevinchern Apr 14, 2026 •

edited

Loading

anahitamansouri Apr 15, 2026 •

edited

Loading

kevinchern Apr 15, 2026 •

edited

Loading