Fix naive branch general stat

tskit-dev · Jul 22, 2024 · d3e17a9 · d3e17a9
1 parent f7f679a
commit d3e17a9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/tests/test_tree_stats.py b/python/tests/test_tree_stats.py
@@ -193,7 +193,7 @@ def naive_branch_general_stat(
     else:
         out = windowed_tree_stat(ts, sigma, windows, span_normalise=span_normalise)
     if drop_time_windows:
-        assert out.shape[1] == 3
+        assert out.ndim == 3
         out = out[:, 0]
     return out