Revert "avoid branching, handle higher-dim case"

MarcoGorelli · MarcoGorelli · commit 06e35b6a634b · 2024-02-15T16:39:25.000Z
This reverts commit 44af268.
diff --git a/sklearn/model_selection/_search.py b/sklearn/model_selection/_search.py
@@ -1079,15 +1079,16 @@ def _store(key_name, array, weights=None, splits=False, rank=False):
                 param_results["param_%s" % name][cand_idx] = value
         for key in param_results:
             arr = np.array(list(param_results[key].values()))
-            # Use one MaskedArray and mask all the places where the param is not
-            # applicable for that candidate (which may not contain all the params).
-            ma = MaskedArray(
-                np.empty((n_candidates, *arr.shape[1:])), mask=True, dtype=arr.dtype
-            )
-            for index, value in param_results[key].items():
-                # Setting the value at an index unmasks that index
-                ma[index] = value
-            param_results[key] = ma
+            if len(arr) == n_candidates:
+                param_results[key] = MaskedArray(arr, mask=False)
+            else:
+                # Use one MaskedArray and mask all the places where the param is not
+                # applicable for that candidate (which may not contain all the params).
+                ma = MaskedArray(np.empty(n_candidates), mask=True, dtype=arr.dtype)
+                for index, value in param_results[key].items():
+                    # Setting the value at an index unmasks that index
+                    ma[index] = value
+                param_results[key] = ma
 
         results.update(param_results)
         # Store a list of param dicts at the key 'params'
