Add block-spin update sampler #44
Conversation
| tricks were employed. Ideally, an adjacency list can be used, however, adjacencies are ragged, | ||
| which makes vectorization inapplicable. | ||
|
|
||
| Block-Gibbs and Block-Metropolis obey detailed balance and are ergodic methods at finite |
jackraymond
left a comment
jackraymond
left a comment
There was a problem hiding this comment.
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.
| def __init__(self, G: Graph, grbm: GRBM, num_reads: int, kind: str): | ||
| super().__init__() | ||
| topology = G.graph.get("family", None) | ||
| if topology not in {"zephyr", "pegasus", "chimera"}: |
There was a problem hiding this comment.
Maybe make partition an optional argument. It seems like we only have this restriction in order to make the partition?
There was a problem hiding this comment.
i.e. make partition an optional argument. If partition is provided don't check for the graph or topology arguments.
A user may provide a subgraph of Zephyr as a zephyr family graph for example (e.g. a Chimera[2m,t] inside a Zephyr[m,t] would permit a 2-coloring).
| h = self.linear[block] | ||
| J = self.quadratic[self.padded_adjacencies_weight[block]] | ||
| effective_field = (xnbr * J.unsqueeze(0)).sum(2) + h | ||
| if self.kind == "metropolis": |
There was a problem hiding this comment.
Can be polite and capitalize Mr Metropolis's and Mr Gibb's names, the least they deserve?
| @torch.compile | ||
| @torch.no_grad | ||
| def step_(self, beta: torch.Tensor): | ||
| """Performs a block-Gibbs update in-place. |
| grbm (GRBM): The Graph-Restricted Boltzmann Machine to sample from. | ||
| """ | ||
|
|
||
| def __init__(self, G: Graph, grbm: GRBM, num_reads: int, kind: str): |
There was a problem hiding this comment.
I think 'kind' is called 'proposal_acceptance_criteria' in SimulatedAnnealingSampler. It took me a while to work out what kind was, and would be good to standardize.
| def get_partition(self, topology): | ||
| partition = defaultdict(list) | ||
| if topology == "zephyr": | ||
| lin2lattice = dnx.zephyr_coordinates(self.G.graph['rows']).linear_to_zephyr |
There was a problem hiding this comment.
Before applying this you should check if you should check the node type property of the graph, it may already be coordinates and then this would fail.
| lin2lattice = dnx.zephyr_coordinates(self.G.graph['rows']).linear_to_zephyr | ||
| crayon = dnx.zephyr_four_color | ||
| elif topology == "pegasus": | ||
| lin2lattice = dnx.pegasus_coordinates(self.G.graph['rows']).linear_to_pegasus |
There was a problem hiding this comment.
pegasus graphs support 3-node types, this transformation may be inappriate.
| lin2lattice = dnx.pegasus_coordinates(self.G.graph['rows']).linear_to_pegasus | ||
| crayon = dnx.pegasus_four_color | ||
| elif topology == "chimera": | ||
| lin2lattice = dnx.chimera_coordinates(self.G.graph['rows']).linear_to_chimera |
There was a problem hiding this comment.
chimera graphs support 2 node types, this is not required if the graph is already coordinated.
| self.idx_to_node = grbm.idx_to_node | ||
| self.partition = self.get_partition(topology) | ||
| self.padded_adjacencies, self.padded_adjacencies_weight = self.get_adjacencies() | ||
| self.x = nn.Parameter(rands((num_reads, grbm.n_nodes)), requires_grad=False) |
There was a problem hiding this comment.
x should be allowed as an input, and the random number should allow a seed when its not provided.
| grbm (GRBM): The Graph-Restricted Boltzmann Machine to sample from. | ||
| """ | ||
|
|
||
| def __init__(self, G: Graph, grbm: GRBM, num_reads: int, kind: str): |
There was a problem hiding this comment.
x should also be allowed as an input, from which num_reads could be inferred (num_reads and x both optional).
4 participants
TODOs:
some low-hanging fruits for optimization:
zerosin methodstep_