Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
54 changes: 5 additions & 49 deletions .devcontainer/devcontainer.json
Original file line number Diff line number Diff line change
@@ -1,51 +1,7 @@
// For format details, see https://aka.ms/devcontainer.json. For config options, see the
// README at: https://github.com/devcontainers/templates/tree/main/src/debian
// For format details, see https://aka.ms/devcontainer.json.
{
"name": "Ocean Development Environment",

// python 3.11 on debian, with latest Ocean and optional packages
// source repo: https://github.com/dwavesystems/ocean-dev-docker
"image": "docker.io/dwavesys/ocean-dev:latest",

// install repo pip requirements (only if present) on content update
"updateContentCommand": "[ ! -r requirements.txt ] || pip install -r requirements.txt",

// forward/expose container services (relevant only when run locally)
"forwardPorts": [
// dwave-inspector web app
18000, 18001, 18002, 18003, 18004,
// OAuth connect redirect URIs
36000, 36001, 36002, 36003, 36004
],

"portsAttributes": {
"18000-18004": {
"label": "D-Wave Problem Inspector",
"requireLocalPort": true
},
"36000-36004": {
"label": "OAuth 2.0 authorization code redirect URI",
"requireLocalPort": true
}
},

// Configure tool-specific properties.
"customizations": {
// Configure properties specific to VS Code.
"vscode": {
// Set *default* container specific settings.json values on container create.
"settings": {
"workbench": {
"editorAssociations": {
"*.md": "vscode.markdown.preview.editor"
},
"startupEditor": "readme"
}
},
"extensions": [
"ms-python.python",
"ms-toolsai.jupyter"
]
}
}
// Debian stable with the second-latest Python, latest Ocean, and optional dev packages.
// Docker image source: https://github.com/dwavesystems/ocean-dev-docker.
// Devcontainer config: https://github.com/dwavesystems/ocean-devcontainer.
"image": "docker.io/dwavesys/ocean-dev:latest"
}
24 changes: 24 additions & 0 deletions .github/PULL_REQUEST_TEMPLATE.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
<!--
Thanks for contributing a pull request! Please ensure that
your PR conforms to our contributor guide.
https://github.com/dwave-examples/template-dash/blob/main/CONTRIBUTING.md
-->

### Associated GitHub Issue
<!--If applicable. Example: Closes https://github.com/dwave-examples/template-dash/issues/X-->

### Feature Implemented/Bugs Fixed
<!--Please explain your changes.-->

### Additional Information
<!--Any additional information you think is important.-->

### Accessibility Score
<!--If the UI was significantly changed, what is the updated accessibility score?-->

### AI Generation Disclosure
<!-- If AI was used in the preparation of this pull request, please disclose
the tool(s) used, how they were used, and specify what code or text is AI generated.
If no AI tools were used, please write "No AI tools used" in this section. Read our
policy on AI generated code at
https://docs.dwavequantum.com/en/latest/ocean/ai_policy.html -->
1 change: 1 addition & 0 deletions requirements.txt
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
dash[diskcache]~=3.2
dash-bootstrap-components~=2.0
dwave-experimental==2026.5.27
dwave-ocean-sdk~=9.0
dwave-pytorch-plugin~=0.3
einops~=0.8
Expand Down
119 changes: 92 additions & 27 deletions src/model_wrapper.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
from typing import Optional

import numpy as np
from dimod.exceptions import BinaryQuadraticModelStructureError
import plotly.express as px
import torch
import yaml
Expand Down Expand Up @@ -306,14 +307,25 @@ def step(self, batch: tuple[torch.Tensor, torch.Tensor], epoch: int) -> torch.Te
self.losses["mse_losses"].append(mse_loss.item())

with torch.no_grad():
samples = self._grbm.sample( # type: ignore
sampler=self.sampler,
prefactor=self.PREFACTOR,
linear_range=self.linear_range,
quadratic_range=self.quadratic_range,
device=spins.device,
sample_params=self.sampler_kwargs,
)
try:
samples = self._grbm.sample( # type: ignore
sampler=self.sampler,
prefactor=self.PREFACTOR,
linear_range=self.linear_range,
quadratic_range=self.quadratic_range,
device=spins.device,
sample_params=self.sampler_kwargs,
)
except BinaryQuadraticModelStructureError:
self._rebuild_sampler()
samples = self._grbm.sample( # type: ignore
sampler=self.sampler,
prefactor=self.PREFACTOR,
linear_range=self.linear_range,
quadratic_range=self.quadratic_range,
device=spins.device,
sample_params=self.sampler_kwargs,
)

spins = spins.reshape(-1, spins.shape[-1])

Expand All @@ -329,17 +341,31 @@ def step(self, batch: tuple[torch.Tensor, torch.Tensor], epoch: int) -> torch.Te
# train boltzmann machine
if train_grbm(self._tpar["opt_step"], epoch):
self._grbm_optimizer.zero_grad()
grbm_loss, self._tpar["sample_set"] = nll_loss(
spins=spins.detach(),
grbm=self._grbm,
sampler=self.sampler,
sampler_kwargs=self.sampler_kwargs,
linear_range=self.linear_range,
quadratic_range=self.quadratic_range,
prefactor=self.PREFACTOR,
persistent_qpu_sample_helper=self._tpar["persistent_qpu_sample_helper"],
sample_set=self._tpar["sample_set"],
)
try:
grbm_loss, self._tpar["sample_set"] = nll_loss(
spins=spins.detach(),
grbm=self._grbm,
sampler=self.sampler,
sampler_kwargs=self.sampler_kwargs,
linear_range=self.linear_range,
quadratic_range=self.quadratic_range,
prefactor=self.PREFACTOR,
persistent_qpu_sample_helper=self._tpar["persistent_qpu_sample_helper"],
sample_set=self._tpar["sample_set"],
)
except BinaryQuadraticModelStructureError:
self._rebuild_sampler()
grbm_loss, self._tpar["sample_set"] = nll_loss(
spins=spins.detach(),
grbm=self._grbm,
sampler=self.sampler,
sampler_kwargs=self.sampler_kwargs,
linear_range=self.linear_range,
quadratic_range=self.quadratic_range,
prefactor=self.PREFACTOR,
persistent_qpu_sample_helper=self._tpar["persistent_qpu_sample_helper"],
sample_set=self._tpar["sample_set"],
)
grbm_loss.backward()
self._grbm_optimizer.step()

Expand All @@ -352,6 +378,34 @@ def step(self, batch: tuple[torch.Tensor, torch.Tensor], epoch: int) -> torch.Te

return mse_loss

def _rebuild_sampler(self) -> None:
"""Reconnect to the QPU and rebuild the sampler after a structure error.

Used as a Zephyr-only fallback when the cached embedding references qubits
or couplers that are no longer available on the QPU. ``EmbeddingComposite``
re-embeds the GRBM's logical graph dynamically onto the current working graph,
so the GRBM weights remain valid and no retraining is needed.
"""
from dwave.system import DWaveSampler, EmbeddingComposite

print(
f"Sampler structure error detected; reconnecting to QPU '{self.qpu}' and rebuilding sampler."
)
qpu_sampler = DWaveSampler(solver=self.qpu)
qpu_topology = qpu_sampler.properties["topology"]["type"]

if qpu_topology != "zephyr":
raise RuntimeError(
f"Sampler rebuild is only supported for Zephyr QPUs (got '{qpu_topology}')."
)

# EmbeddingComposite finds a valid minor embedding of the GRBM's logical
# graph on the current working QPU topology without requiring retraining.
self.sampler = EmbeddingComposite(qpu_sampler)
self.linear_range = qpu_sampler.properties["h_range"]
self.quadratic_range = qpu_sampler.properties["j_range"]
print(f"Sampler rebuilt with EmbeddingComposite for Zephyr QPU '{self.qpu}'.")

def generate_output(self, latent_qpu_file: str, sharpen: bool = False, save_to_file: str = "") -> go.Figure:
"""Generate output images from trained model.

Expand All @@ -366,14 +420,25 @@ def generate_output(self, latent_qpu_file: str, sharpen: bool = False, save_to_f
self._grbm.eval()

with torch.no_grad():
samples = self._grbm.sample(
self.sampler,
prefactor=self.PREFACTOR,
device=self._device,
linear_range=self.linear_range,
quadratic_range=self.quadratic_range,
sample_params=self.sampler_kwargs,
)
try:
samples = self._grbm.sample(
self.sampler,
prefactor=self.PREFACTOR,
device=self._device,
linear_range=self.linear_range,
quadratic_range=self.quadratic_range,
sample_params=self.sampler_kwargs,
)
except BinaryQuadraticModelStructureError:
self._rebuild_sampler()
samples = self._grbm.sample(
self.sampler,
prefactor=self.PREFACTOR,
device=self._device,
linear_range=self.linear_range,
quadratic_range=self.quadratic_range,
sample_params=self.sampler_kwargs,
)

with open(latent_qpu_file, "w") as f:
json.dump(samples[0].tolist(), f)
Expand Down
21 changes: 16 additions & 5 deletions src/utils/callback_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@
from torchvision.utils import save_image

from demo_configs import GENERATE_NEW_MODEL_DIAGRAM, GRAPH_COLORS, SHARPEN_OUTPUT, THEME_COLOR_SECONDARY
from src.utils.common import get_graph_mapping, greedy_get_subgraph
from src.utils.common import get_graph_mapping, greedy_get_subgraph, _try_zephyr_sublattice_fallback

MODEL_PATH = Path("models")
JSON_FILE_DIR = "generated_json"
Expand Down Expand Up @@ -358,13 +358,24 @@ def generate_model_fig(
"""
qpu = DWaveSampler(solver=qpu)
qpu_graph = qpu.to_networkx_graph()
subgraph = greedy_get_subgraph(n_nodes=n_latents, random_seed=random_seed, graph=qpu_graph)
_, mapping = get_graph_mapping(subgraph)
qpu_topology = qpu.properties["topology"]["type"]

latent_mapping = [mapping[node] for node in subgraph.nodes()]
try:
subgraph = greedy_get_subgraph(n_nodes=n_latents, random_seed=random_seed, graph=qpu_graph)
_, mapping = get_graph_mapping(subgraph)
latent_mapping = [mapping[node] for node in subgraph.nodes()]
except Exception:
if qpu_topology != "zephyr":
raise
fallback = _try_zephyr_sublattice_fallback(
n_nodes=n_latents, qpu_graph=qpu_graph, random_seed=random_seed
)
if fallback is None:
raise
subgraph, _ = fallback
latent_mapping = list(subgraph.nodes())

qpu_shape = qpu.properties["topology"]["shape"][0]
qpu_topology = qpu.properties["topology"]["type"]

if qpu_topology == "pegasus":
node_coords = dnx.drawing.pegasus_layout(dnx.pegasus_graph(qpu_shape), crosses=True)
Expand Down
Loading