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.

dwave/plugins/torch/samplers/bipartite_sampler.py

tests/test_bipartite_sampler.py

dwave/plugins/torch/samplers/bipartite_sampler.py

tests/test_bipartite_sampler.py

kevinchern

Added a couple minor requests

dwave/plugins/torch/samplers/bipartite_sampler.py

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

dwave/plugins/torch/samplers/bipartite_sampler.py

+            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

dwave/plugins/torch/samplers/bipartite_sampler.py

+            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

dwave/plugins/torch/samplers/bipartite_sampler.py

+            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

dwave/plugins/torch/samplers/bipartite_sampler.py

+
+        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

dwave/plugins/torch/samplers/bipartite_sampler.py

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


Suggested change

return mask

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

dwave/plugins/torch/samplers/bipartite_sampler.py

+
+        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

dwave/plugins/torch/samplers/block_spin_sampler.py

+            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)

dwave/plugins/torch/samplers/block_spin_sampler.py

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.

dwave/plugins/torch/samplers/bipartite_sampler.py

kevinchern

Almost there!

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

tests/test_block_sampler.py

+        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-23T18:59:31Z

tests/test_block_sampler.py

+
+        sampler = BlockSampler(grbm, crayon, num_chains=2, schedule=[1.0])
+
+        # Case 1: Valid single block unclamped


Can each of these test cases be wrapped in individual subtests?

releasenotes/notes/conditional-sampling-94c0d541df3e3fa8.yaml

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

dwave/plugins/torch/samplers/block_spin_sampler.py

+            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

tests/test_block_sampler.py

+        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

tests/test_bipartite_sampler.py

+        # 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

tests/test_bipartite_sampler.py

+    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.

tests/test_bipartite_sampler.py

dwave/plugins/torch/samplers/bipartite_sampler.py

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

releasenotes/notes/conditional-sampling-94c0d541df3e3fa8.yaml

+    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

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

	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"]


		sampler = BlockSampler(grbm, crayon, num_chains=2, schedule=[1.0])

		# Case 1: Valid single block unclamped

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

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!

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!

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

kevinchern left a comment •

edited

Loading

anahitamansouri Mar 24, 2026 •

edited

Loading