Fix monotone constraint with tuple input. (dmlc#7891)

trivialfis · Aug 12, 2022 · 46d9392 · 46d9392
1 parent a55d3bd
commit 46d9392
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/python-package/xgboost/core.py b/python-package/xgboost/core.py
@@ -1403,10 +1403,12 @@ def __init__(
         self.set_param(params_processed or {})
 
     def _transform_monotone_constrains(
-        self, value: Union[Dict[str, int], str]
+        self, value: Union[Dict[str, int], str, Tuple[int, ...]]
     ) -> Union[Tuple[int, ...], str]:
         if isinstance(value, str):
             return value
+        if isinstance(value, tuple):
+            return value
 
         constrained_features = set(value.keys())
         feature_names = self.feature_names or []

diff --git a/tests/python/test_monotone_constraints.py b/tests/python/test_monotone_constraints.py
@@ -93,6 +93,11 @@ def test_monotone_constraints(self, tree_method: str, policy: str) -> None:
         constrained = xgb.train(params_for_constrained, training_dset)
         assert is_correctly_constrained(constrained)
 
+    def test_monotone_constraints_tuple(self) -> None:
+        params_for_constrained = {"monotone_constraints": (1, -1)}
+        constrained = xgb.train(params_for_constrained, training_dset)
+        assert is_correctly_constrained(constrained)
+
     @pytest.mark.parametrize('format', [dict, list])
     def test_monotone_constraints_feature_names(self, format):