from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple from contextlib import contextmanager, nullcontext from dataclasses import dataclass import torch from gem.flow_matching.noise_distribution import NoiseDistribution from gem.flow_matching import flow_matching_utils # Re-export energy utilities for external use from .sampler_energy import ( energy_and_grads_batch, energy_batch, register_property_time_override, ) # Proposals from .proposals.base import Proposal from .proposals.random_proposal import RandomProposal from .proposals.gwd_proposal import GWDProposal from .proposals.gwg_block_proposal import GWGBlockProposal from .proposals.dlangevin_proposal import ( DLangevinProposal, DLangevinVectorizedProposal, DLangevinNoMHProposal, DLangevinMTProposal, DLangevinAnnealingProposal, DLangevinTwoBetasProposal, DLangevinTwoBetasVectorizedProposal, DLangevinTwoBetasAnnealingProposal, DLangevinTwoBetasAnnealingVectorizedProposal, DLangevinTwoBetasAnnealingVectorizedNoOriginProposal, ) from .proposals.simple_proposal import ( SimpleProposal, SimpleProposalV2, SimpleProposalV2GuidedStrong, run_simple_v2_warmup_vectorized, ) from .dlangevin_utils import ( TWO_BETA_PROPOSALS, TWO_BETA_SCALAR_PROPOSALS, TWO_BETA_VECTOR_PROPOSALS, TWO_BETA_ANNEALING_PROPOSALS, TWO_BETA_ANNEALING_VECTOR_PROPOSALS, ) _VECTORIZED_SIMPLE_WARMUP_PROPOSALS = frozenset({"simple_ver2", "simple_v2"}) @dataclass(frozen=True) class MCMCStepEvent: """Read-only views of an MCMC step for diagnostics.""" step: int nodes: Sequence[torch.Tensor] edges: Sequence[torch.Tensor] energies: torch.Tensor accepted: torch.Tensor proposed_nodes: Sequence[torch.Tensor] proposed_edges: Sequence[torch.Tensor] proposed_energies: Optional[torch.Tensor] def _amp_autocast_ctx(device: torch.device, amp_dtype: Optional[str] = None): """Autocast context controlled by desired AMP dtype. If `amp_dtype` is None or empty, returns a no-op context to keep FP32. Otherwise, uses the requested dtype ("fp16"/"float16" or "bf16"/"bfloat16"). """ if device.type != "cuda": return nullcontext() if not amp_dtype: return nullcontext() amp = amp_dtype.lower() if amp in {"bf16", "bfloat16"}: return torch.amp.autocast("cuda", dtype=torch.bfloat16) if amp in {"fp16", "float16"}: return torch.amp.autocast("cuda", dtype=torch.float16) return nullcontext() def should_vectorize_simple_warmup( proposal: str, *, vectorized: bool = True, ) -> bool: """Return whether a configured warmup can use the dense simple-v2 runner.""" return bool(vectorized) and str(proposal or "").strip().lower() in ( _VECTORIZED_SIMPLE_WARMUP_PROPOSALS ) def make_proposal(method: str, **kwargs) -> Proposal: """Factory for proposal mechanisms.""" m = (method or "").lower() if m == "random": return RandomProposal() if m == "gwd": beta = float(kwargs.get("gwd_beta", 1.0)) return GWDProposal(beta=beta) if m in {"gwg", "gwg_block", "gibbs_grad"}: beta = float(kwargs.get("gwd_beta", 1.0)) return GWGBlockProposal(beta=beta) if m in {"dlangevin", "dlang", "dl"}: beta = float(kwargs.get("dl_beta", 1.0)) lambda_X = float(kwargs.get("dl_lambda_X", 1.0)) lambda_E = float(kwargs.get("dl_lambda_E", 1.0)) return DLangevinProposal( beta=beta, lambda_X=lambda_X, lambda_E=lambda_E, ) if m in {"dlangevin_vec", "dlangevin_vectorized", "dlang_vec", "dl_vec"}: beta = float(kwargs.get("dl_beta", 1.0)) lambda_X = float(kwargs.get("dl_lambda_X", 1.0)) lambda_E = float(kwargs.get("dl_lambda_E", 1.0)) return DLangevinVectorizedProposal( beta=beta, lambda_X=lambda_X, lambda_E=lambda_E, ) if m in {"dlangevin_mt", "dlangevinmt", "dlang_mt", "dl_mt", "dlangevin_multi"}: beta = float(kwargs.get("dl_beta", 1.0)) lambda_X = float(kwargs.get("dl_lambda_X", 1.0)) lambda_E = float(kwargs.get("dl_lambda_E", 1.0)) num_tries_raw = ( kwargs.get("dl_num_tries", None) or kwargs.get("dl_mt_k", None) or kwargs.get("dl_k", None) ) energy_batch_size_raw = ( kwargs.get("dl_mt_energy_batch", None) or kwargs.get("dl_energy_chunk", None) ) try: num_tries = int(num_tries_raw) if num_tries_raw is not None else 1 except (TypeError, ValueError): num_tries = 1 try: energy_batch_size = int(energy_batch_size_raw) if energy_batch_size_raw is not None else 256 except (TypeError, ValueError): energy_batch_size = 256 return DLangevinMTProposal( beta=beta, lambda_X=lambda_X, lambda_E=lambda_E, num_tries=num_tries, energy_batch_size=energy_batch_size, ) if m in TWO_BETA_SCALAR_PROPOSALS: beta_prop = kwargs.get("dl_beta_prop") beta_mh = kwargs.get("dl_beta_mh") if beta_prop is None or beta_mh is None: raise ValueError("dlangevintwobetas requires dl_beta_prop and dl_beta_mh to be set.") lambda_X = float(kwargs.get("dl_lambda_X", 1.0)) lambda_E = float(kwargs.get("dl_lambda_E", 1.0)) return DLangevinTwoBetasProposal( beta_prop=float(beta_prop), beta_mh=float(beta_mh), lambda_X=lambda_X, lambda_E=lambda_E, ) if m in TWO_BETA_VECTOR_PROPOSALS: beta_prop = kwargs.get("dl_beta_prop") beta_mh = kwargs.get("dl_beta_mh") if beta_prop is None or beta_mh is None: raise ValueError( "dlangevin_two_betas_vec requires dl_beta_prop and dl_beta_mh to be set." ) lambda_X = float(kwargs.get("dl_lambda_X", 1.0)) lambda_E = float(kwargs.get("dl_lambda_E", 1.0)) return DLangevinTwoBetasVectorizedProposal( beta_prop=float(beta_prop), beta_mh=float(beta_mh), lambda_X=lambda_X, lambda_E=lambda_E, ) if m in {"dlangevin_annealing", "dlang_annealing", "dl_annealing"}: beta_init = float(kwargs.get("dl_beta_init", kwargs.get("dl_beta", 1.0))) beta_final = float(kwargs.get("dl_beta_final", beta_init)) anneal_steps_raw = kwargs.get("dl_beta_anneal_steps", kwargs.get("dl_anneal_steps", 0)) try: anneal_steps = max(int(anneal_steps_raw), 0) except (TypeError, ValueError): anneal_steps = 0 lambda_X = float(kwargs.get("dl_lambda_X", 1.0)) lambda_E = float(kwargs.get("dl_lambda_E", 1.0)) return DLangevinAnnealingProposal( beta_init=beta_init, beta_final=beta_final, anneal_steps=anneal_steps, lambda_X=lambda_X, lambda_E=lambda_E, ) if m in TWO_BETA_ANNEALING_PROPOSALS: beta_prop = kwargs.get("dl_beta_prop") beta_mh_init = kwargs.get("dl_beta_mh_init") beta_mh_final = kwargs.get("dl_beta_mh_final") beta_mh_anneal_steps = kwargs.get("dl_beta_mh_anneal_steps") missing = [ k for k, v in [ ("dl_beta_prop", beta_prop), ("dl_beta_mh_init", beta_mh_init), ("dl_beta_mh_final", beta_mh_final), ("dl_beta_mh_anneal_steps", beta_mh_anneal_steps), ] if v is None ] if missing: raise ValueError(f"dlangevintwobetas_annealing requires {', '.join(missing)}.") lambda_X = float(kwargs.get("dl_lambda_X", 1.0)) lambda_E = float(kwargs.get("dl_lambda_E", 1.0)) if m == "dlangevin_two_betas_annealing_vec_no_origin": proposal_cls = DLangevinTwoBetasAnnealingVectorizedNoOriginProposal elif m in TWO_BETA_ANNEALING_VECTOR_PROPOSALS: proposal_cls = DLangevinTwoBetasAnnealingVectorizedProposal else: proposal_cls = DLangevinTwoBetasAnnealingProposal return proposal_cls( beta_prop=float(beta_prop), beta_mh_init=float(beta_mh_init), beta_mh_final=float(beta_mh_final), beta_mh_anneal_steps=int(beta_mh_anneal_steps), lambda_X=lambda_X, lambda_E=lambda_E, ) if m in {"dlangevin_nomh", "dlang_nomh", "dl_nomh"}: beta = float(kwargs.get("dl_beta", 1.0)) lambda_X = float(kwargs.get("dl_lambda_X", 1.0)) lambda_E = float(kwargs.get("dl_lambda_E", 1.0)) return DLangevinNoMHProposal( beta=beta, lambda_X=lambda_X, lambda_E=lambda_E, ) if m == "simple": edits = kwargs.get("simple_n_edits", 5) edits_int = 5 if edits is None else int(edits) return SimpleProposal(edits_per_step=edits_int, use_conditioner=True) if m in {"simple_ver2", "simple_v2"}: edits = kwargs.get("simple_n_edits", 5) edits_int = 5 if edits is None else int(edits) return SimpleProposalV2(edits_per_step=edits_int, use_conditioner=False) if m in {"simple_ver2_guided", "simple_v2_guided"}: edits = kwargs.get("simple_n_edits", 5) edits_int = 5 if edits is None else int(edits) return SimpleProposalV2(edits_per_step=edits_int, use_conditioner=True) if m in {"simple_ver2_guided_strong", "simple_v2_guided_strong"}: edits = kwargs.get("simple_n_edits", 5) edits_int = 5 if edits is None else int(edits) return SimpleProposalV2GuidedStrong(edits_per_step=edits_int) raise ValueError( f"Unknown proposal '{method}'. Supported: ['random', 'gwd', 'gwg_block', 'dlangevin', 'dlangevin_vec', 'dlangevin_mt', 'dlangevin_annealing', 'dlangevintwobetas', 'dlangevin_two_betas_vec', 'dlangevintwobetas_annealing', 'dlangevin_two_betas_annealing_vec', 'dlangevin_two_betas_annealing_vec_no_origin', 'dlangevin_noMH', 'simple', 'simple_ver2', 'simple_ver2_guided', 'simple_ver2_guided_strong']." ) def initialize_random_graphs( batch_size: int, dataset_info, device: torch.device = torch.device("cpu"), transition: str = "marginal", ) -> List[Tuple[torch.Tensor, torch.Tensor]]: """Sample a batch of random graphs used as MCMC initial states.""" # Sample number of nodes from empirical distribution n_nodes = dataset_info.nodes_dist.sample_n(batch_size, device) n_max = torch.max(n_nodes).item() arange = torch.arange(n_max, device=device).unsqueeze(0).expand(batch_size, -1) node_mask = arange < n_nodes.unsqueeze(1) # Sample node and edge types from reference noise distribution noise_dist = NoiseDistribution(transition, dataset_info) limit_dist = noise_dist.get_limit_dist() z_T = flow_matching_utils.sample_discrete_feature_noise( limit_dist=limit_dist, node_mask=node_mask ) graphs: List[Tuple[torch.Tensor, torch.Tensor]] = [] for i in range(batch_size): n = int(n_nodes[i].item()) node_types = torch.argmax(z_T.X[i, :n], dim=-1) edge_types = torch.argmax(z_T.E[i, :n, :n], dim=-1) graphs.append((node_types.cpu(), edge_types.cpu())) return graphs @contextmanager def _model_eval_ctx(model: torch.nn.Module): """ Temporarily set the model to eval() and restore the previous training mode. Used to disable dropout for the entire MCMC section. """ prev = model.training try: if prev: model.eval() yield finally: if prev: model.train() def mcmc_sample_batch( model, dataset_info, node_types_list: Sequence[torch.Tensor], edge_types_list: Sequence[torch.Tensor], extra_features, domain_features, steps: int = 10, device: torch.device = torch.device("cpu"), proposal: str = "random", gwd_beta: float = 1.0, dl_beta: float = 1.0, dl_beta_init: Optional[float] = None, dl_beta_final: Optional[float] = None, dl_beta_anneal_steps: Optional[int] = None, dl_beta_prop: Optional[float] = None, dl_beta_mh: Optional[float] = None, dl_beta_mh_init: Optional[float] = None, dl_beta_mh_final: Optional[float] = None, dl_beta_mh_anneal_steps: Optional[int] = None, dl_lambda_X: float = 1.0, dl_lambda_E: float = 1.0, dl_num_tries: Optional[int] = None, dl_mt_energy_batch: Optional[int] = None, simple_n_edits: Optional[int] = None, energy_split_threshold: Optional[float] = None, dl_beta_near: Optional[float] = None, dl_lambda_X_near: Optional[float] = None, dl_lambda_E_near: Optional[float] = None, dl_beta_far: Optional[float] = None, dl_lambda_X_far: Optional[float] = None, dl_lambda_E_far: Optional[float] = None, amp_dtype: Optional[str] = None, collect_stats: bool = False, step_offset: int = 0, property_lambda_prop: Optional[float] = None, property_t_override: Optional[float] = None, property_t_schedule: Optional[Dict[str, Any]] = None, step_callback: Optional[Callable[[MCMCStepEvent], None]] = None, ): """Run parallel MCMC chains (one per graph) starting from the provided batch. Returns ------- tuple (node_types_list, edge_types_list, n_accept_total, n_steps_total) Notes on performance: - Chain states are moved to `device` once here and kept device-resident throughout the loop. We only move them back to CPU on return. - Model is put into eval() for the entire MCMC to disable dropout. - If provided, `step_callback` is called after each state update with device-resident accepted and proposed states. Callers must clone data they retain. """ assert len(node_types_list) == len(edge_types_list) B = len(node_types_list) if B == 0: return [], [], 0, 0, {} try: step_offset_int = int(step_offset) except (TypeError, ValueError): step_offset_int = 0 # --- Move chain states to device ONCE, keep them there during the loop --- node_types_list = [t.to(device, non_blocking=True) for t in node_types_list] edge_types_list = [t.to(device, non_blocking=True) for t in edge_types_list] # Keep a copy of chain origins for distance-to-origin stats if collect_stats: origin_nodes = [t.clone() for t in node_types_list] origin_edges = [t.clone() for t in edge_types_list] else: origin_nodes = [] origin_edges = [] dual_params = ( energy_split_threshold is not None and dl_beta_near is not None and dl_lambda_X_near is not None and dl_lambda_E_near is not None and dl_beta_far is not None and dl_lambda_X_far is not None and dl_lambda_E_far is not None ) if dual_params and proposal.lower() not in {"dlangevin", "dlang", "dl"}: raise ValueError( "Dual DLangevin parameters require proposal='dlangevin'." ) if dual_params and step_callback is not None: raise ValueError( "step_callback is not supported with dual DLangevin parameters." ) if proposal.lower() in TWO_BETA_PROPOSALS: if dl_beta_prop is None or dl_beta_mh is None: raise ValueError( f"proposal='{proposal}' requires dl_beta_prop and dl_beta_mh." ) if proposal.lower() in TWO_BETA_ANNEALING_PROPOSALS: missing = [ k for k, v in [ ("dl_beta_prop", dl_beta_prop), ("dl_beta_mh_init", dl_beta_mh_init), ("dl_beta_mh_final", dl_beta_mh_final), ("dl_beta_mh_anneal_steps", dl_beta_mh_anneal_steps), ] if v is None ] if missing: raise ValueError(f"dlangevintwobetas_annealing requires {', '.join(missing)}.") total_accepts = 0 # Stats accumulators for single-parameter mode stats: Dict[str, Any] = {} if collect_stats and not dual_params: stats.update( dict( total_proposals=0, total_accepted=0, nontriv_any=0, nontriv_node=0, nontriv_edge=0, acc_nontriv_any=0, acc_nontriv_node=0, acc_nontriv_edge=0, # Distances to origin accumulators (proposed and accepted) prop_dist_nodes_sum=0.0, prop_dist_edges_sum=0.0, acc_dist_nodes_sum=0.0, acc_dist_edges_sum=0.0, # Step size accumulators (current -> proposed) step_prop_nodes_sum=0.0, step_prop_edges_sum=0.0, step_acc_nodes_sum=0.0, step_acc_edges_sum=0.0, step_acc_total_sq_sum=0.0, ) ) track_distance = collect_stats total_distance_nodes = 0.0 total_distance_edges = 0.0 B_initial = B prop_impl: Optional[Proposal] = None if not dual_params: beta_init = dl_beta if dl_beta_init is None else float(dl_beta_init) beta_final = dl_beta if dl_beta_final is None else float(dl_beta_final) try: anneal_steps_val = None if dl_beta_anneal_steps is None else int(dl_beta_anneal_steps) except (TypeError, ValueError): anneal_steps_val = None prop_impl = make_proposal( proposal, gwd_beta=gwd_beta, dl_beta=dl_beta, dl_beta_init=beta_init, dl_beta_final=beta_final, dl_beta_anneal_steps=anneal_steps_val, dl_beta_prop=dl_beta_prop, dl_beta_mh=dl_beta_mh, dl_beta_mh_init=dl_beta_mh_init, dl_beta_mh_final=dl_beta_mh_final, dl_beta_mh_anneal_steps=dl_beta_mh_anneal_steps, dl_lambda_X=dl_lambda_X, dl_lambda_E=dl_lambda_E, dl_num_tries=dl_num_tries, dl_mt_energy_batch=dl_mt_energy_batch, simple_n_edits=simple_n_edits, ) if ( property_lambda_prop is not None and prop_impl is not None and bool(getattr(prop_impl, "supports_property_lambda_override", False)) ): prop_impl.set_property_lambda_override(property_lambda_prop) def _resolve_property_t(step_idx: int) -> Optional[float]: if property_t_override is not None: try: return float(property_t_override) except (TypeError, ValueError): return None if not property_t_schedule: return None mode = str(property_t_schedule.get("mode", "linear")).strip().lower() if mode not in {"linear", "lin"}: return None try: steps_total = int(property_t_schedule.get("steps", 0) or 0) t_start = float(property_t_schedule.get("t_start", 0.0)) t_end = float(property_t_schedule.get("t_end", 1.0)) except (TypeError, ValueError): return None if steps_total <= 1: return t_end max_idx = max(steps_total - 1, 1) s = min(max(int(step_idx), 0), max_idx) ratio = float(s) / float(max_idx) return float(t_start + (t_end - t_start) * ratio) def _set_property_t(step_idx: int) -> None: register_property_time_override(_resolve_property_t(step_idx)) # Disable dropout for the entire MCMC run (restored on exit) with _model_eval_ctx(model): try: _set_property_t(0) # Current energies for all graphs (detached for acceptance decisions) # Use autocast only if requested by caller with _amp_autocast_ctx(device, amp_dtype): current_E = energy_batch( model=model, node_types_list=node_types_list, edge_types_list=edge_types_list, dataset_info=dataset_info, device=device, extra_features=extra_features, domain_features=domain_features, detach=True, ) # (B,) if not dual_params: assert prop_impl is not None for step_idx in range(steps): _set_property_t(step_idx) prop_impl.on_step_start(step_offset_int + step_idx) prev_nodes = [t.clone() for t in node_types_list] if track_distance else None prev_edges = [t.clone() for t in edge_types_list] if track_distance else None # === Propose === prop_result = prop_impl.propose( model=model, dataset_info=dataset_info, node_types_list=node_types_list, edge_types_list=edge_types_list, extra_features=extra_features, domain_features=domain_features, device=device, amp_dtype=amp_dtype, ) # === Score proposals if needed by the proposal === prop_E: Optional[torch.Tensor] = None if prop_impl.needs_proposed_energy(): with _amp_autocast_ctx(device, amp_dtype): prop_E = energy_batch( model=model, node_types_list=prop_result.prop_nodes, edge_types_list=prop_result.prop_edges, dataset_info=dataset_info, device=device, extra_features=extra_features, domain_features=domain_features, detach=True, ) # (B,) # === Accept/Reject (early-reject + caching inside) === accept_mask, prop_E_from_accept = prop_impl.accept( model=model, dataset_info=dataset_info, current_nodes=node_types_list, current_edges=edge_types_list, prop_result=prop_result, current_E=current_E, prop_E=prop_E, extra_features=extra_features, domain_features=domain_features, device=device, amp_dtype=amp_dtype, ) # ---- Stats accumulation (use pre-update current state) ---- if collect_stats: B_local = len(node_types_list) stats["total_proposals"] += B_local # Loop per-chain due to variable sizes for b in range(B_local): nt_cur = node_types_list[b] et_cur = edge_types_list[b] nt_prop = prop_result.prop_nodes[b] et_prop = prop_result.prop_edges[b] # Step change (current -> proposed) # Nodes h_node_step = int((nt_prop != nt_cur).sum().item()) # Edges: count only i 0) nontriv_edge = (h_edge_step > 0) nontriv_any = (nontriv_node or nontriv_edge) stats["nontriv_node"] += int(nontriv_node) stats["nontriv_edge"] += int(nontriv_edge) stats["nontriv_any"] += int(nontriv_any) acc_b = bool(accept_mask[b].item()) if acc_b: stats["total_accepted"] += 1 if acc_b and nontriv_any: stats["acc_nontriv_any"] += 1 if acc_b and nontriv_node: stats["acc_nontriv_node"] += 1 if acc_b and nontriv_edge: stats["acc_nontriv_edge"] += 1 # Accumulate step sizes (current -> proposed) stats["step_prop_nodes_sum"] += float(h_node_step) stats["step_prop_edges_sum"] += float(h_edge_step) if acc_b: stats["step_acc_nodes_sum"] += float(h_node_step) stats["step_acc_edges_sum"] += float(h_edge_step) stats["step_acc_total_sq_sum"] += float( (h_node_step + h_edge_step) ** 2 ) # Distances to origin nt0 = origin_nodes[b] et0 = origin_edges[b] prop_dist_nodes = int((nt_prop != nt0).sum().item()) if et_prop.numel() == 0: prop_dist_edges = 0 else: diff0 = et_prop != et0 prop_dist_edges = int(torch.triu(diff0, diagonal=1).sum().item()) stats["prop_dist_nodes_sum"] += float(prop_dist_nodes) stats["prop_dist_edges_sum"] += float(prop_dist_edges) if acc_b: stats["acc_dist_nodes_sum"] += float(prop_dist_nodes) stats["acc_dist_edges_sum"] += float(prop_dist_edges) # Update states and energies where accepted (keep everything on device) # Ensure mask device matches energies for indexing accept_mask_dev = accept_mask.to(current_E.device) n_acc = int(accept_mask.sum().item()) total_accepts += n_acc if n_acc > 0: acc_idx = torch.nonzero(accept_mask, as_tuple=False).flatten().tolist() for i in acc_idx: node_types_list[i] = prop_result.prop_nodes[i] edge_types_list[i] = prop_result.prop_edges[i] eff_prop_E = prop_E if prop_E is not None else prop_E_from_accept assert eff_prop_E is not None, "Proposal did not provide proposed energies." current_E[accept_mask_dev] = eff_prop_E[accept_mask_dev] if track_distance and prev_nodes is not None and prev_edges is not None: for prev_nt, prev_et, cur_nt, cur_et in zip(prev_nodes, prev_edges, node_types_list, edge_types_list): dist_nodes = int((cur_nt != prev_nt).sum().item()) if cur_et.numel() == 0: dist_edges = 0 else: diff_step = cur_et != prev_et dist_edges = int(torch.triu(diff_step, diagonal=1).sum().item()) total_distance_nodes += float(dist_nodes) total_distance_edges += float(dist_edges) if step_callback is not None: step_callback( MCMCStepEvent( step=step_offset_int + step_idx, nodes=node_types_list, edges=edge_types_list, energies=current_E, accepted=accept_mask, proposed_nodes=prop_result.prop_nodes, proposed_edges=prop_result.prop_edges, proposed_energies=( prop_E if prop_E is not None else prop_E_from_accept ), ) ) else: threshold_val = float(energy_split_threshold) def _run_subset( mask: torch.Tensor, beta_val: float, lambda_X_val: float, lambda_E_val: float, ) -> int: nonlocal current_E, node_types_list, edge_types_list if mask is None: return 0 idx = torch.nonzero(mask, as_tuple=False).flatten() if idx.numel() == 0: return 0 idx_list = idx.tolist() sub_nodes = [node_types_list[i] for i in idx_list] sub_edges = [edge_types_list[i] for i in idx_list] prop_local = DLangevinProposal( beta=beta_val, lambda_X=lambda_X_val, lambda_E=lambda_E_val, ) prop_res = prop_local.propose( model=model, dataset_info=dataset_info, node_types_list=sub_nodes, edge_types_list=sub_edges, extra_features=extra_features, domain_features=domain_features, device=device, amp_dtype=amp_dtype, ) prop_E_local: Optional[torch.Tensor] = None if prop_local.needs_proposed_energy(): with _amp_autocast_ctx(device, amp_dtype): prop_E_local = energy_batch( model=model, node_types_list=prop_res.prop_nodes, edge_types_list=prop_res.prop_edges, dataset_info=dataset_info, device=device, extra_features=extra_features, domain_features=domain_features, detach=True, ) sub_idx = idx.to(current_E.device) current_E_subset = current_E.index_select(0, sub_idx) accept_mask_local, prop_E_updated = prop_local.accept( model=model, dataset_info=dataset_info, current_nodes=sub_nodes, current_edges=sub_edges, prop_result=prop_res, current_E=current_E_subset, prop_E=prop_E_local, extra_features=extra_features, domain_features=domain_features, device=device, amp_dtype=amp_dtype, ) eff_prop_E_local = prop_E_local if prop_E_local is not None else prop_E_updated assert eff_prop_E_local is not None, "Proposal did not provide updated energies." acc_idx_local = torch.nonzero(accept_mask_local, as_tuple=False).flatten().tolist() for rel_idx in acc_idx_local: global_idx = idx_list[rel_idx] node_types_list[global_idx] = prop_res.prop_nodes[rel_idx] edge_types_list[global_idx] = prop_res.prop_edges[rel_idx] current_E.index_copy_(0, sub_idx, eff_prop_E_local) return int(accept_mask_local.sum().item()) for step_idx in range(steps): _set_property_t(step_idx) prev_nodes = [t.clone() for t in node_types_list] if track_distance else None prev_edges = [t.clone() for t in edge_types_list] if track_distance else None mask_near = current_E < threshold_val mask_far = ~mask_near total_accepts += _run_subset( mask_near, float(dl_beta_near), float(dl_lambda_X_near), float(dl_lambda_E_near), ) total_accepts += _run_subset( mask_far, float(dl_beta_far), float(dl_lambda_X_far), float(dl_lambda_E_far), ) if track_distance and prev_nodes is not None and prev_edges is not None: for prev_nt, prev_et, cur_nt, cur_et in zip(prev_nodes, prev_edges, node_types_list, edge_types_list): dist_nodes = int((cur_nt != prev_nt).sum().item()) if cur_et.numel() == 0: dist_edges = 0 else: diff_step = cur_et != prev_et dist_edges = int(torch.triu(diff_step, diagonal=1).sum().item()) total_distance_nodes += float(dist_nodes) total_distance_edges += float(dist_edges) finally: register_property_time_override(None) # Return CPU tensors to keep external API stable # Finalize stats means if collect_stats and stats.get("total_proposals", 0) > 0: total_props = max(int(stats["total_proposals"]), 1) total_acc = max(int(stats["total_accepted"]), 0) n_any = max(int(stats["nontriv_any"]), 0) n_node = max(int(stats["nontriv_node"]), 0) n_edge = max(int(stats["nontriv_edge"]), 0) mean_step_acc_total = ( (stats["step_acc_nodes_sum"] + stats["step_acc_edges_sum"]) / max(total_acc, 1) if total_acc > 0 else 0.0 ) variance_step_acc_total = ( max( stats["step_acc_total_sq_sum"] / total_acc - mean_step_acc_total**2, 0.0, ) if total_acc > 0 else 0.0 ) stats.update( dict( overall_accept=(stats["total_accepted"] / total_props), accept_nontrivial_any=(stats["acc_nontriv_any"] / n_any) if n_any > 0 else 0.0, accept_nontrivial_node=(stats["acc_nontriv_node"] / n_node) if n_node > 0 else 0.0, accept_nontrivial_edge=(stats["acc_nontriv_edge"] / n_edge) if n_edge > 0 else 0.0, mean_prop_distance_nodes=(stats["prop_dist_nodes_sum"] / total_props), mean_prop_distance_edges=(stats["prop_dist_edges_sum"] / total_props), mean_prop_distance_total=((stats["prop_dist_nodes_sum"] + stats["prop_dist_edges_sum"]) / total_props), mean_acc_distance_nodes=(stats["acc_dist_nodes_sum"] / max(total_acc, 1)) if total_acc > 0 else 0.0, mean_acc_distance_edges=(stats["acc_dist_edges_sum"] / max(total_acc, 1)) if total_acc > 0 else 0.0, mean_acc_distance_total=((stats["acc_dist_nodes_sum"] + stats["acc_dist_edges_sum"]) / max(total_acc, 1)) if total_acc > 0 else 0.0, # Mean step sizes per proposal and per accepted mean_step_distance_nodes=(stats["step_prop_nodes_sum"] / total_props), mean_step_distance_edges=(stats["step_prop_edges_sum"] / total_props), mean_step_distance_total=((stats["step_prop_nodes_sum"] + stats["step_prop_edges_sum"]) / total_props), mean_step_acc_distance_nodes=(stats["step_acc_nodes_sum"] / max(total_acc, 1)) if total_acc > 0 else 0.0, mean_step_acc_distance_edges=(stats["step_acc_edges_sum"] / max(total_acc, 1)) if total_acc > 0 else 0.0, mean_step_acc_distance_total=mean_step_acc_total, std_step_acc_distance_total=variance_step_acc_total**0.5, ) ) distance_total_nodes = float(total_distance_nodes) distance_total_edges = float(total_distance_edges) distance_total = distance_total_nodes + distance_total_edges steps_effective = max(int(steps), 1) chains_effective = max(int(B_initial), 1) distance_per_step = distance_total / steps_effective if steps_effective > 0 else 0.0 denom = steps_effective * chains_effective distance_per_chain_step = distance_total / denom if denom > 0 else 0.0 distance_summary = dict( distance_total=distance_total, distance_total_nodes=distance_total_nodes, distance_total_edges=distance_total_edges, distance_per_step=distance_per_step, distance_per_chain_step=distance_per_chain_step, ) stats.update(distance_summary) node_types_cpu = [t.detach().cpu() for t in node_types_list] edge_types_cpu = [t.detach().cpu() for t in edge_types_list] return node_types_cpu, edge_types_cpu, total_accepts, steps * B, (stats if collect_stats else {})