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.

kevinchern · 2026-04-07T17:02:21Z

        super().__init__()

    def sample(self, x: torch.Tensor | None = None) -> torch.Tensor:
        """Sample from the dimod sampler and return the corresponding tensor.


Update docstrings to reflect behaviour of conditional sampling

kevinchern · 2026-04-07T17:06:03Z

+        for i in range(x.shape[0]):
+            # Fresh BQM
+            bqm = dimod.BinaryQuadraticModel.from_ising(h, J)
+
+            # Build conditioning dict
+            conditioned = {node: int(x[i, j].item())
+                        for j, node in enumerate(nodes) if mask[i, j]}


Can we iterate over x instead of range(x.shape[0])? Seems more natural

kevinchern · 2026-04-07T17:13:06Z

+                results.append(full)
+                continue
+
+            # Sample one configuration per input


This needs to be documented in the docstring. It should also raise a warning if num_reads is supplied and overwritten.

kevinchern · 2026-04-07T17:14:11Z

+            # Sample one configuration per input
+            sample_kwargs = dict(self._sampler_params)
+            sample_kwargs["num_reads"] = 1
+            sample_set = AggregatedSamples.spread(


What's the motivation for this line?

Oh you mean calling the spread? I was following the earlier use case. You're right, it's redundant. I'll remove it then.

kevinchern · 2026-04-07T17:17:36Z

+            sample = sample_set.first.sample
+
+            # Reconstruct full sample
+            full = torch.empty(n_nodes, device=device)
+            for j, node in enumerate(nodes):
+                full[j] = conditioned[node] if node in conditioned else float(sample[node])
+            results.append(full)


sample's variable ordering may not be identical to that of grbm.nodes.
I'd add a test to verify correct ordering first.

@anahitamansouri missed this one?

No I didn't. The order is the same as that. As you suggested I have written the test to show that. You can see your sampleset_to_tensor also uses this order.

kevinchern · 2026-04-07T17:18:59Z

@@ -22,7 +22,7 @@
 from dwave.plugins.torch.samplers.dimod_sampler import DimodSampler


Add test to verify correct ordering of variables

kevinchern · 2026-04-15T22:54:42Z


        The sample set returned from the latest sample call is stored in :func:`DimodSampler.sample_set`
-        which is overwritten by subsequent calls.
+        which is overwritten by subsequent calls. When ``x`` is provided (conditional sampling), exactly


I'm wondering if there's a better way to document and handle this. The sample size parameter of a sampler is not guaranteed to use num_reads. @anahitamansouri any thoughts?

Yeah, this was also my challenge.

kevinchern · 2026-04-15T22:58:23Z

+                    "(one sample per input row).",
+                    UserWarning,
+                )
+                sample_kwargs["num_reads"] = 1


@anahitamansouri do you think this is sufficient? i.e., can we safely ignore the num_reads parameters altogether and change the documentation to ~= "only the first sample is retained ..."

Alternatively, and perhaps the right way to do this, is to return tensors of shape $(b, N, d)$ where $b$ is batch size, $N$ is the sample size (i.e., what we enforce to be $1$ here), and $d$ is model dimension / number of qubits.

Agreed. We discussed to keep it (b, d) to be consistent with other samplers as you remember. I will update the documentation.

kevinchern · 2026-04-15T23:00:43Z

+        self._sample_set = dimod.SampleSet.from_samples(
+            (samples.cpu().numpy(), self._grbm.nodes),
+            vartype=dimod.SPIN,
+            energy=energies,
+        )


Why explicitly reconstruct the sample set?

Setting self._sample_set part is only necessary to be consistent with your code as you have this field in the class and sample_set function needs that. It is not necessary for the logic of sample function. If you remember, I didn't have this in the earlier version and I told you I have not set this. This is one way I though of addressing it. Do you mean this way of creating samples is wrong? Do you have any suggestions? :)

Hmm... I think there are at least a couple simplifications, 1) self._sample_set = sample_set and note in documentation that it's the latest sample set, 2) concatenate all samplesets

kevinchern · 2026-04-15T23:02:52Z

+            sample = sample_set.first.sample
+
+            # Reconstruct full sample
+            full = torch.empty(n_nodes, device=device)
+            for j, node in enumerate(nodes):
+                full[j] = conditioned[node] if node in conditioned else float(sample[node])
+            results.append(full)


@anahitamansouri missed this one?

kevinchern · 2026-04-15T23:34:00Z

+        ])
+
+        out = sampler.sample(x)
+        ss = sampler.sample_set


edit: A test that fails this

# Unconditional sampling if x is None: self._sample_set = AggregatedSamples.spread( self._sampler.sample_ising(h, J, **self._sampler_params) ) return sampleset_to_tensor(nodes, self._sample_set, device)

kevinchern · 2026-04-16T17:45:14Z

+
+            sample_set = self._sampler.sample(bqm, **sample_kwargs)
+
+            # Extract sampled values


Raise an error here if sample set has sample size > 1

kevinchern · 2026-04-16T17:59:42Z

+
+            # Reconstruct full sample
+            full = torch.empty(n_nodes, device=device)
+            for j, node in enumerate(nodes):


Suggested change

for j, node in enumerate(nodes):

for node, idx in grbm._node_to_index.items():

kevinchern · 2026-04-16T18:08:25Z

+            self._sample_set = AggregatedSamples.spread(
+                self._sampler.sample_ising(h, J, **self._sampler_params)
+            )
+            return sampleset_to_tensor(nodes, self._sample_set, device)


We should not use sampleset_to_tensor because that function is only safe when contained in GRBM class.
Here, we should rely on grbm._index_to_node and grbm._node_to_index etc. to guarantee correct mapping between indices and nodes

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

Updating the node order

95e5f1d

		@@ -22,7 +22,7 @@
		from dwave.plugins.torch.samplers.dimod_sampler import DimodSampler


		sample_set = self._sampler.sample(bqm, **sample_kwargs)

		# Extract sampled values

	for j, node in enumerate(nodes):
	for node, idx in grbm._node_to_index.items():

Conversation

anahitamansouri commented Apr 7, 2026

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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

anahitamansouri Apr 16, 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!

anahitamansouri Apr 16, 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 Apr 16, 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!

kevinchern Apr 15, 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 Apr 16, 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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

anahitamansouri Apr 16, 2026 •

edited

Loading

anahitamansouri Apr 16, 2026 •

edited

Loading

kevinchern Apr 16, 2026 •

edited

Loading

kevinchern Apr 15, 2026 •

edited

Loading

kevinchern Apr 16, 2026 •

edited

Loading