Coerce NODE_IS_SAMPLE to correct numpy type before negation

Fixes #937
tskit-dev · Jul 17, 2024 · 2155335 · 2155335
1 parent 342454c
commit 2155335
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/tsinfer/inference.py b/tsinfer/inference.py
@@ -1970,7 +1970,10 @@ def get_samples_tree_sequence(self, map_additional_sites=True):
 
         # All true ancestors are samples in the ancestors tree sequence. We unset
         # the SAMPLE flag but keep other flags intact.
-        new_flags = np.bitwise_and(tables.nodes.flags, ~tskit.NODE_IS_SAMPLE)
+        new_flags = tables.nodes.flags
+        new_flags = np.bitwise_and(
+            new_flags, ~new_flags.dtype.type(tskit.NODE_IS_SAMPLE)
+        )
         tables.nodes.flags = new_flags.astype(np.uint32)
         sample_ids = list(self.sample_id_map.values())
         assert len(tables.nodes) == sample_ids[0]