Updating Dimod Sampler by anahitamansouri · Pull Request #73 · dwavesystems/dwave-pytorch-plugin

anahitamansouri · 2026-04-07T16:20:15Z

This PR adds:

Conditional sampling feature for Dimod sampler.
Tests for this feature.

kevinchern

Did a quick first pass over the main function.

thisac · 2026-06-04T17:29:04Z

        """
-        if x is not None:
-            raise NotImplementedError("Support for conditional sampling has not been implemented.")
+        device = self._grbm._linear.device


Might not be an issue, but accessing hidden attributes (as in GRBM._linear here) likely indicates that it either shouldn't be hidden or accessed in a different way (e.g., GRBM.linear if that exists).

Noted. I also use _linear inconsistently and should use the public attribute

I recommend changing all other similar occurrences in this PR, e.g., self._sampler_params

I updated access to private attributes of external objects ( GRBM._linear and GRBM._node_to_idx) but self._sampler_params is an internal attribute of DimodSampler, so I kept that unchanged.

kevinchern

Lookin' goooood! Left some minor comments

kevinchern · 2026-06-10T18:09:53Z

+            torch.Tensor: 
+            - If ``x is None``:
+                tensor of shape ``(num_reads, n_nodes)``.
+
+            - If ``x`` is provided:
+                tensor of shape ``(batch_size, num_reads, n_nodes)``.


@thisac want to double check with you whether there's a conventional format for docstrings where return value shapes are input-dependent

I would just write it as a single paragraph (no specific formatting). As long as it renders nicely in the docs, should be OK. Bullets are fine, but not sure how the indentations after each bullet are rendered.

kevinchern · 2026-06-10T18:11:25Z

+            raise ValueError(f"x must have shape (batch_size, {n_nodes})")
+
+        mask = ~torch.isnan(x)
+        if not torch.all(torch.isin(x[mask], torch.tensor([-1, 1], device=x.device))):


Use consistently device or x.device

kevinchern · 2026-06-10T18:12:18Z

-        if x is not None:
-            raise NotImplementedError("Support for conditional sampling has not been implemented.")
+        device = self._grbm._linear.device
+        nodes = self._grbm.nodes


Looks like this is only ever used once; no need to define this variable here

kevinchern · 2026-06-10T22:33:51Z

+                for v in bqm.linear:
+                    if bqm.linear[v] > ub:
+                        bqm.set_linear(v, ub)
+                    elif bqm.linear[v] < lb:
+                        bqm.set_linear(v, lb)
+
+            # Clip quadratic biases
+            if self._quadratic_range is not None:
+                lb, ub = self._quadratic_range
+                for u, v, bias in bqm.iter_quadratic():
+                    if bias > ub:
+                        bqm.set_quadratic(u, v, ub)
+                    elif bias < lb:
+                        bqm.set_quadratic(u, v, lb)


Why do the two for-loops differ? Here it's for v in bqm.linear and the next uses iter_quadratic

bqm.linear is a mapping from variable to bias, while bqm.iter_quadratic() gives (u, v, bias) triples. Are you suggesting to change the linear loop to bqm.linear.items() (style of the iteration) or do you think the loops can be merged? Could you clarify this? :)

Oops, I meant why not use bqm.iter_linear()?

kevinchern · 2026-06-10T22:36:57Z

+                    elif bias < lb:
+                        bqm.set_quadratic(u, v, lb)
+
+            # Sample one configuration per input


This comment is dated

kevinchern · 2026-06-10T22:52:55Z

+
+                free_values = samples[:, i, :][free_mask]
+                self.assertTrue(torch.all(torch.isin(free_values, torch.tensor([-1.0, 1.0]))),
+                                f"Free variables should be sampled as ±1 in read {i}") 


Suggested change

f"Free variables should be sampled as ±1 in read {i}")

f"Free variables should be sampled as ±1 in read {i}")

Should all subtests be separated by a blank line?

kevinchern · 2026-06-10T22:54:28Z

+                    x,
+                    msg=f"Fully clamped inputs should be preserved in read {i}"
+                ) 
+        with self.subTest("Conditional sampling supports mixed fully-clamped and partially-clamped rows."):


Suggested change

with self.subTest("Conditional sampling supports mixed fully-clamped and partially-clamped rows."):

with self.subTest("Conditional sampling supports mixed fully clamped and partially clamped rows."):

kevinchern · 2026-06-10T22:57:05Z

+
+            # Sample one configuration per input
+            sample_kwargs = dict(self._sampler_params)
+


Suggested change

(empty space characters; use formatter)

thisac · 2026-06-11T21:20:43Z

+            torch.Tensor: 
+            - If ``x is None``:
+                tensor of shape ``(num_reads, n_nodes)``.
+
+            - If ``x`` is provided:
+                tensor of shape ``(batch_size, num_reads, n_nodes)``.


I would just write it as a single paragraph (no specific formatting). As long as it renders nicely in the docs, should be OK. Bullets are fine, but not sure how the indentations after each bullet are rendered.

thisac · 2026-06-12T00:03:08Z

+        samples = torch.stack(results, dim=0)
+        return samples
+
+    def _sampleset_to_tensor(self, sample_set: SampleSet, device: Optional[torch.device] = None


This seems close to identical to utils.sampleset_to_tensor(). Why is this added here and used instead of the one in utils?

If you follow our discussion above, you can see that Kevin mentionedutils.sampleset_to_tensor should bot be used here because that function is only safe when contained in GRBM class.
Here, I defined another function that rely on grbm._index_to_node and grbm._node_to_index order. So, there is a subtle difference between these functions.

Maybe I'm missing something, but it's not clear to me why this is different to

def _sampleset_to_tensor(...) ordered_vars = [v for v, _ in sorted(self._grbm.node_to_idx.items(), key=lambda x: x[1])] return utils.sampleset_to_tensor(ordered_vars, sample_set)

Seems to me still much better to use what we already have defined in utils.

We should not use sampleset_to_tensor because that function is only safe when contained in GRBM class.

@kevinchern Not sure I understand what this means or why. If it's very specifically just for GRBM classes, but we have similar needs in other places, we should probably either generalize it or move it into the GRBM scope.

Here, we should rely on grbm._index_to_node and grbm._node_to_index etc. to guarantee correct mapping between indices and nodes

@kevinchern Isn't that what we're doing when passing in ordered_vars?

The difference is that in this new function we follow the grbm.node_to_idx order whereas in the one in utils we allow any ordered_vars and in the older version in DimodSampler we passed the grbm.nodes as ordered_vars. The idea is that grbm.nodes could have a different order from grbm.node_to_idx and we want to make sure we follow grbm.node_to_idx order.

kevinchern

LGTM. Can merge after @thisac approves!

thisac

A couple of tests should be added for coverage. Otherwise looks good to me.

Happy to approve as soon as these have been addressed!

thisac · 2026-07-14T18:07:06Z

 from __future__ import annotations

-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Optional


Minor, but I'd recommend using ... | None instead of Optional[...] below.

Suggested change

from typing import TYPE_CHECKING, Any, Optional

from typing import TYPE_CHECKING, Any

thisac · 2026-07-14T18:08:31Z

    import dimod
-
-    from dwave.plugins.torch.models.boltzmann_machine import GraphRestrictedBoltzmannMachine
+    from dimod import SampleSet


Minor, but since you're already importing dimod and also use e.g., dimod.Sampler for typing, it might be a bit cleaner to just use dimod.SampleSet instead of importing SampleSet again here.

Suggested change

from dimod import SampleSet

thisac · 2026-07-14T18:23:46Z

+        if x.shape[1] != n_nodes:
+            raise ValueError(f"x must have shape (batch_size, {n_nodes})")
+
+        mask = ~torch.isnan(x)
+        if not torch.all(torch.isin(x[mask], torch.tensor([-1, 1], device=device))):
+            raise ValueError("x must contain only ±1 or NaN")


Neither of these cases are tested.

thisac · 2026-07-14T18:26:41Z

+            # Clip linear biases for remaining free variables
+            if self._linear_range is not None:
+                lb, ub = self._linear_range
+                for v, bias in bqm.iter_linear():
+                    if bias > ub:
+                        bqm.set_linear(v, ub)
+                    elif bias < lb:
+                        bqm.set_linear(v, lb)
+
+            # Clip quadratic biases
+            if self._quadratic_range is not None:
+                lb, ub = self._quadratic_range
+                for u, v, bias in bqm.iter_quadratic():
+                    if bias > ub:
+                        bqm.set_quadratic(u, v, ub)
+                    elif bias < lb:
+                        bqm.set_quadratic(u, v, lb)


It doesn't seem like there's any test coverage for either of these either. Could you check whether there are any tests that jump into these blocks?

Thanks for this suggestion. I added the test for the linear biases as requested. For the quadratic ones, I noticed that the quadratic clipping is already applied in to_ising() before constructing the BQM. Unlike linear biases, conditioning cannot make quadratic biases exceed their specified range; therefore, the conditional quadratic clipping loop will not encounter out-of-range quadratic biases in the current implementation. This makes the block redundant, so I will remove it.

thisac · 2026-07-14T18:26:55Z

+
+        reference_shape = results[0].shape
+        if not all(result.shape == reference_shape for result in results):
+            raise ValueError(f"Expected all samples to have shape {reference_shape}")


Also missing test coverage.

Thanks for pointing this out. This check is intended to guard against inconsistent sample shapes returned by a sampler. Under normal execution, I cannot trigger this error. And this is why I added a small fake sampler that intentionally returns a different number of reads for different conditional rows. This simulates a sampler returning inconsistent sample sizes and verifies that DimodSampler raises the expected ValueError instead of failing later during stacking.

Updating Dimod Sampler

21bc015

anahitamansouri requested a review from kevinchern April 7, 2026 16:20

anahitamansouri self-assigned this Apr 7, 2026

anahitamansouri added the enhancement New feature or request label Apr 7, 2026

kevinchern requested changes Apr 7, 2026

View reviewed changes

add test, release note and improve sample

0cf8a89

kevinchern requested changes Apr 15, 2026

View reviewed changes

kevinchern reviewed Apr 16, 2026

View reviewed changes

Comment thread dwave/plugins/torch/samplers/dimod_sampler.py Outdated

kevinchern reviewed Apr 16, 2026

View reviewed changes

Comment thread dwave/plugins/torch/samplers/dimod_sampler.py Outdated

kevinchern reviewed Apr 16, 2026

View reviewed changes

Comment thread dwave/plugins/torch/samplers/dimod_sampler.py Outdated

amansouribigvand added 2 commits April 28, 2026 16:17

Updating the node order

95e5f1d

support multiple reads

923fdac

anahitamansouri requested a review from kevinchern June 3, 2026 22:42

thisac reviewed Jun 4, 2026

View reviewed changes

kevinchern requested changes Jun 10, 2026

View reviewed changes

thisac reviewed Jun 12, 2026

View reviewed changes

formatting, adding some checks, etc

416a8a0

anahitamansouri requested review from kevinchern and thisac June 19, 2026 15:11

kevinchern approved these changes Jul 13, 2026

View reviewed changes

thisac reviewed Jul 14, 2026

View reviewed changes

Adding tests

a0c2b10

	f"Free variables should be sampled as ±1 in read {i}")
	f"Free variables should be sampled as ±1 in read {i}")

	with self.subTest("Conditional sampling supports mixed fully-clamped and partially-clamped rows."):
	with self.subTest("Conditional sampling supports mixed fully clamped and partially clamped rows."):


		# Sample one configuration per input
		sample_kwargs = dict(self._sampler_params)

	from typing import TYPE_CHECKING, Any, Optional
	from typing import TYPE_CHECKING, Any

Uh oh!

Conversation

anahitamansouri commented Apr 7, 2026

Uh oh!

kevinchern left a comment

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!

Uh oh!

Uh oh!

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!

Choose a reason for hiding this comment

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!

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!

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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

anahitamansouri Jun 17, 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!

anahitamansouri Jun 19, 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!

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

anahitamansouri Jun 17, 2026 •

edited

Loading

anahitamansouri Jun 19, 2026 •

edited

Loading

kevinchern left a comment •

edited

Loading

anahitamansouri Aug 3, 2026 •

edited

Loading