From a3e7187c4f180084921e2f68dc867196acf13d1d Mon Sep 17 00:00:00 2001 From: Andrew Rowley Date: Thu, 10 Oct 2024 11:09:53 +0100 Subject: [PATCH] Fix these calls too --- mcmc/mcmc_coordinator_vertex.py | 7 +++---- mcmc/mcmc_vertex.py | 4 ++-- 2 files changed, 5 insertions(+), 6 deletions(-) diff --git a/mcmc/mcmc_coordinator_vertex.py b/mcmc/mcmc_coordinator_vertex.py index 04e3bf6..e16320a 100644 --- a/mcmc/mcmc_coordinator_vertex.py +++ b/mcmc/mcmc_coordinator_vertex.py @@ -163,7 +163,7 @@ def get_sequence_mask(self, placement, routing_info): def get_acknowledge_key(self, placement, routing_info): if self._is_receiver_placement(placement): - key = routing_info.get_first_key_from_pre_vertex( + key = routing_info.get_safe_first_key_from_pre_vertex( placement.vertex, self._acknowledge_partition_name) return key return 0 @@ -224,10 +224,9 @@ def generate_data_specification( mcmc_placement = FecDataView.get_placement_of_vertex(vertex) self._mcmc_placements.append(mcmc_placement) if self._is_receiver_placement(mcmc_placement): - key = routing_info.get_first_key_from_pre_vertex( + key = routing_info.get_safe_first_key_from_pre_vertex( vertex, self._acknowledge_partition_name) - if key is not None: - keys.append(key) + keys.append(key) keys.sort() # Write the data size in words diff --git a/mcmc/mcmc_vertex.py b/mcmc/mcmc_vertex.py index e4c4503..ee15f84 100644 --- a/mcmc/mcmc_vertex.py +++ b/mcmc/mcmc_vertex.py @@ -147,7 +147,7 @@ def _get_model_state_array(self): def get_result_key(self, placement): if self._is_receiver_placement(placement): routing_info = FecDataView.get_routing_infos() - key = routing_info.get_first_key_from_pre_vertex( + key = routing_info.get_safe_first_key_from_pre_vertex( placement.vertex, self._result_partition_name) return key return 0 @@ -163,7 +163,7 @@ def _is_receiver_placement(self, placement): def get_cholesky_result_key(self, placement): if self._is_cholesky_receiver_placement(placement): routing_info = FecDataView.get_routing_infos() - key = routing_info.get_first_key_from_pre_vertex( + key = routing_info.get_safe_first_key_from_pre_vertex( placement.vertex, self._cholesky_result_partition_name) return key return 0