Fix these calls too

SpiNNakerManchester · Oct 10, 2024 · a3e7187 · a3e7187
1 parent 6104452
commit a3e7187
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 6 deletions.
diff --git 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
@@ -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