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:50:28Z

+            self.assertTupleEqual(samples.shape, (2, 1, 4))
+            samples = samples.squeeze()
+
+            # Check clamped values unchanged
+            mask = ~torch.isnan(x)
+            self.assertTrue(torch.all(samples[mask] == x[mask]))
+
+            # Check free variables are ±1
+            free_mask = torch.isnan(x)
+            free_values = samples[free_mask]
+            self.assertTrue(torch.all(torch.isin(free_values, torch.tensor([-1.0, 1.0]))),
+                            "Free variables should be sampled as ±1")


Should these be separate subtests? They're currently part of a subtest but only the second test is appropriate for that heading

I agree. I wanted to be concise :) but it makes sense to have multiple subtests.

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)

kevinchern · 2026-06-10T22:59:42Z

+                    full_read[:, idx] = torch.from_numpy(sample_array[:, var_to_idx[node]]).to(device=device, dtype=torch.float)
+            results.append(full_read)
+        # Stack to get (batch_size, num_reads, n_nodes)
+        samples = torch.stack(results, dim=0)


Unlikely but certainly possible uncaught error here: results have tensors of different sample size (e.g., num_reads) and so the stack function will fail. Probably want to check all elements in result have same shape.

I thought this would not happen because I am appending tensors of shape(num_reads, num_nodes) to results and then stack them. So, they would have the same shape. I will add this check anyways.

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.

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

	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)

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!

Choose a reason for hiding this comment

Uh oh!

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

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!

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

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

anahitamansouri Jun 11, 2026 •

edited

Loading

anahitamansouri Jun 18, 2026 •

edited

Loading

anahitamansouri Jun 17, 2026 •

edited

Loading

anahitamansouri Jun 19, 2026 •

edited

Loading