Introduce special case for upcoming numpy release 2.0.0

const-ae · Apr 16, 2024 · eb385d9 · eb385d9
1 parent b3a163f
commit eb385d9
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/pylemur/tl/_design_matrix_utils.py b/src/pylemur/tl/_design_matrix_utils.py
@@ -1,6 +1,7 @@
 from collections.abc import Mapping
 
 import numpy as np
+from numpy.lib import NumpyVersion
 import pandas as pd
 
 # import patsy
@@ -74,6 +75,8 @@ def convert_formula_to_design_matrix(formula, obs_data):
 
 def row_groups(matrix, return_reduced_matrix=False, return_group_ids=False):
     reduced_matrix, inv = np.unique(matrix, axis=0, return_inverse=True)
+    if NumpyVersion(np.__version__) >= '2.0.0rc':
+        inv = np.squeeze(inv)
     group_ids = np.unique(inv)
     if return_reduced_matrix and return_group_ids:
         return inv, reduced_matrix, group_ids