Add block-spin update sampler by kevinchern · Pull Request #44 · dwavesystems/dwave-pytorch-plugin

kevinchern · 2025-11-14T21:56:53Z

This PR introduces a block-spin update sampler.
It includes duplicate code from #40

dwave/plugins/torch/samplers/block_spin_update.py

jackraymond

Great job, I'll come back to this when you add tests. Per your summary. Thanks for getting this draft up quickly.

Would be nice to have a tidy public GPU code to use and cite for speedup relative to our CPU implementation in higher throughput applications as soon as possible.

dwave/plugins/torch/samplers/block_spin_update.py

kevinchern · 2025-12-10T16:31:37Z

Samplers should have a consistent interface. Related: #52, #47

kevinchern · 2025-12-11T16:18:45Z

@jackraymond I updated the PR:

added tests (some notable mentions below)
- test gibbs update by checking magnetization is close to expected
- test metropolis update by checking magnetization is close to expected
- test metropolis update by checking oscillatory behavior for beta = 0 case
added documentation (first pass)
capitalized Metropolis/Gibbs and also make it case-insensitive
The sampler has a simplified signature. It no longer tries to do the colouring itself. Instead, it requests for a colouring function. This generalizes it beyond CPZ topologies. One could, for example, default to a colouring heuristic when no colouring is given, but I've opted away from that for modularity.
~~5. Now a metadata dictionary is returned along with the sample.~~

VolodyaCO

I left some comments about form.

Overall looks fantastic. Thank you.

dwave/plugins/torch/nn/functional.py

dwave/plugins/torch/samplers/block_spin_sampler.py

dwave/plugins/torch/tensor.py

dwave/plugins/torch/utils.py

VolodyaCO · 2025-12-15T19:14:48Z

Can the user compile a method of an existing nn.Module subclassed instance? Ideally, we would want the user to decide whether they want to compile the _step method. It seems as though torch.compile(module) compiles the __call__ method, so maybe we can wrap _step in it and time the compiled and non-compiled versions.

kevinchern · 2025-12-15T23:34:03Z

Addressed @thisac and @VolodyaCO 's PRs. Briefly,

Remove inheritance of nn.Module
Add functional module docstring
Add explicit named args for, and renamed to randspin
Improved documentation and type-hinting

dwave/plugins/torch/samplers/block_spin_sampler.py

jackraymond

Looks good tests ran for me.

I'd recommend you allow setting of x as an initial condition. Check it doesn't move when num_sweeps is none, and check it oscillates under BlockMetropolis when temperature is infinity as tests perhaps.

kevinchern · 2025-12-17T18:31:48Z

Looks good tests ran for me.

I'd recommend you allow setting of x as an initial condition. Check it doesn't move when num_sweeps is none, and check it oscillates under BlockMetropolis when temperature is infinity as tests perhaps.

In the latest commit, I made the following changes:

Fixed documentation RE nonzero & finite temperature regimes.
Allow setting of initial states.
- Test by checking sampler.x.tolist() is equal to input initial states list
Added custom errors to distinguish from ValueError

@jackraymond @thisac

VolodyaCO

All comments have been addressed. Thank you @kevinchern

dwave/plugins/torch/samplers/block_spin_sampler.py

thisac

Thanks @kevinchern! 🎄

dwave/plugins/torch/nn/functional.py

dwave/plugins/torch/samplers/block_spin_sampler.py

dwave/plugins/torch/tensor.py

dwave/plugins/torch/utils.py

thisac · 2025-12-23T23:38:37Z

tests/test_block_sampler.py

+    ZEPHYR = dnx.zephyr_graph(1, coordinates=True)
+    GRBM_ZEPHYR = GRBM(ZEPHYR.nodes, ZEPHYR.edges)
+    CRAYON_ZEPHYR = dnx.zephyr_four_color
+
+    BIPARTITE = nx.complete_bipartite_graph(5, 3)
+    GRBM_BIPARTITE = GRBM(BIPARTITE.nodes, BIPARTITE.edges)
+    def CRAYON_BIPARTITE(b): return b < 5
+
+    GRBM_SINGLE = GRBM([0], [])
+    def CRAYON_SINGLE(s): 0
+
+    GRBM_CRAYON_TEST_CASES = [(GRBM_ZEPHYR, CRAYON_ZEPHYR),
+                              (GRBM_BIPARTITE, CRAYON_BIPARTITE),
+                              (GRBM_SINGLE, CRAYON_SINGLE)]


Better put in a setUpClass() method.

If I understood correctly, that would be incompatible with @parameterized.expand

tests/test_block_spin_sampler.py

tests/test_block_sampler.py

jackraymond · 2026-01-07T05:19:11Z

dwave/plugins/torch/samplers/block_spin_sampler.py

+        self._x = nn.Parameter(initial_states.float(), requires_grad=False)
+        self._zeros = nn.Parameter(torch.zeros((num_chains, 1)), requires_grad=False)
+
+    def to(self, device: DeviceLikeType) -> BlockSpinSampler:


Encourage my laziness by setting device to torch.device("cuda") by default?

jackraymond

No changes necessary

Test ran fine for me, got it working for my benchmarking straightfowardly.

So far as I can tell on an L4:
156 ns per spin update for single sample on Zephyr[m=12] problem.
drops to ~ 156/58 ns per spin update with 58 fold parallelism on an L4. Which is good.

thisac

Apart from a few minor comments, LGTM.

dwave/plugins/torch/tensor.py

dwave/plugins/torch/samplers/block_spin_sampler.py

kevinchern requested a review from jackraymond November 14, 2025 21:57

kevinchern marked this pull request as draft November 17, 2025 21:24

kevinchern assigned VolodyaCO and mhramani and unassigned VolodyaCO and mhramani Nov 17, 2025

jackraymond reviewed Nov 27, 2025

View reviewed changes

dwave/plugins/torch/samplers/block_spin_update.py Outdated Show resolved Hide resolved

jackraymond suggested changes Nov 27, 2025

View reviewed changes

jackraymond reviewed Nov 27, 2025

View reviewed changes

dwave/plugins/torch/samplers/block_spin_update.py Outdated Show resolved Hide resolved

jackraymond reviewed Nov 27, 2025

View reviewed changes

dwave/plugins/torch/samplers/block_spin_update.py Outdated Show resolved Hide resolved

kevinchern requested a review from thisac December 11, 2025 16:18

kevinchern marked this pull request as ready for review December 11, 2025 16:19

kevinchern requested a review from jackraymond December 11, 2025 16:19

VolodyaCO reviewed Dec 15, 2025

View reviewed changes

jackraymond reviewed Dec 16, 2025

View reviewed changes

dwave/plugins/torch/samplers/block_spin_sampler.py Outdated Show resolved Hide resolved

jackraymond reviewed Dec 16, 2025

View reviewed changes

dwave/plugins/torch/samplers/block_spin_sampler.py Outdated Show resolved Hide resolved

jackraymond reviewed Dec 17, 2025

View reviewed changes

dwave/plugins/torch/samplers/block_spin_sampler.py Outdated Show resolved Hide resolved

jackraymond suggested changes Dec 17, 2025

View reviewed changes

kevinchern requested a review from jackraymond December 17, 2025 18:32

VolodyaCO approved these changes Dec 22, 2025

View reviewed changes

dwave/plugins/torch/samplers/block_spin_sampler.py Show resolved Hide resolved

thisac reviewed Dec 23, 2025

View reviewed changes

kevinchern requested a review from thisac January 5, 2026 17:29

jackraymond reviewed Jan 7, 2026

View reviewed changes

jackraymond approved these changes Jan 7, 2026

View reviewed changes

thisac approved these changes Jan 8, 2026

View reviewed changes

Add block-spin update sampler

74666ca

kevinchern force-pushed the feature/block-spin branch from 1d67e64 to 74666ca Compare January 8, 2026 15:49

kevinchern merged commit d08b203 into dwavesystems:main Jan 8, 2026
15 checks passed

kevinchern deleted the feature/block-spin branch January 8, 2026 16:15

Conversation

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

Uh oh!

Uh oh!

Uh oh!

jackraymond 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!

kevinchern commented Dec 10, 2025

Uh oh!

kevinchern commented Dec 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

VolodyaCO 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!

VolodyaCO commented Dec 15, 2025

Uh oh!

kevinchern commented Dec 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jackraymond left a comment

Choose a reason for hiding this comment

Uh oh!

kevinchern commented Dec 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

VolodyaCO left a comment

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!

Uh oh!

Uh oh!

thisac Dec 23, 2025

Choose a reason for hiding this comment

Uh oh!

kevinchern Dec 24, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

jackraymond Jan 7, 2026

Choose a reason for hiding this comment

Uh oh!

jackraymond 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!

thisac left a comment

Choose a reason for hiding this comment

kevinchern commented Nov 14, 2025 •

edited

Loading

kevinchern commented Dec 11, 2025 •

edited

Loading

kevinchern commented Dec 15, 2025 •

edited

Loading

kevinchern commented Dec 17, 2025 •

edited

Loading

jackraymond left a comment •

edited

Loading