From 2eadab6297c77a7a151a91bb3bbcb65ed336b8f6 Mon Sep 17 00:00:00 2001
From: adobles96 <alejandro.d.96@gmail.com>
Date: Fri, 19 Apr 2024 18:41:50 -0700
Subject: [PATCH 1/2] Add device support for XGB tuning

---
 torch_frame/gbdt/tuned_xgboost.py | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/torch_frame/gbdt/tuned_xgboost.py b/torch_frame/gbdt/tuned_xgboost.py
index 9b939d32..edd970ac 100644
--- a/torch_frame/gbdt/tuned_xgboost.py
+++ b/torch_frame/gbdt/tuned_xgboost.py
@@ -93,6 +93,7 @@ def objective(
         dtrain: Any,  # xgboost.DMatrix
         dvalid: Any,  # xgboost.DMatrix
         num_boost_round: int,
+        device: str = 'cpu'
     ) -> float:
         r"""Objective function to be optimized.
 
@@ -101,6 +102,9 @@ def objective(
             dtrain (xgboost.DMatrix): Train data.
             dvalid (xgboost.DMatrix): Validation data.
             num_boost_round (int): Number of boosting round.
+            device (str): The device for XGBoost to train on. One of `cpu`,
+                `cuda`, `cuda<ordinal>`, `gpu`, `gpu:<ordinal>`. See XGBoost
+                documentation for details.
 
         Returns:
             float: Best objective value. Root mean squared error for
@@ -118,8 +122,12 @@ def objective(
              else trial.suggest_float('lambda', 1e-8, 1e2, log=True)),
             "alpha":
             (0.0 if not trial.suggest_categorical('use_alpha', [True, False])
-             else trial.suggest_float('alpha', 1e-8, 1e2, log=True))
+             else trial.suggest_float('alpha', 1e-8, 1e2, log=True)),
+            "device": device
         }
+        if device.startswith("gpu") or device.startswith("cuda"):
+            self.params["tree_method"] = "hist"
+
         if self.params["booster"] == "gbtree" or self.params[
                 "booster"] == "dart":
             self.params["max_depth"] = trial.suggest_int("max_depth", 3, 11)
@@ -188,6 +196,7 @@ def _tune(
         tf_val: TensorFrame,
         num_trials: int,
         num_boost_round: int = 2000,
+        device: str = 'cpu'
     ):
         import optuna
         import xgboost
@@ -207,8 +216,8 @@ def _tune(
                                  feature_types=val_feat_type,
                                  enable_categorical=True)
         study.optimize(
-            lambda trial: self.objective(trial, dtrain, dvalid, num_boost_round
-                                         ), num_trials)
+            lambda trial: self.objective(trial, dtrain, dvalid,
+                                         num_boost_round, device), num_trials)
         self.params.update(study.best_params)
 
         self.model = xgboost.train(self.params, dtrain,

From cdba71aa91c63c908aaf9ee88a224ee634257abc Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Sat, 20 Apr 2024 03:17:41 +0000
Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 torch_frame/gbdt/tuned_xgboost.py | 27 +++++++++++----------------
 1 file changed, 11 insertions(+), 16 deletions(-)

diff --git a/torch_frame/gbdt/tuned_xgboost.py b/torch_frame/gbdt/tuned_xgboost.py
index edd970ac..562c53f9 100644
--- a/torch_frame/gbdt/tuned_xgboost.py
+++ b/torch_frame/gbdt/tuned_xgboost.py
@@ -88,13 +88,12 @@ def _to_xgboost_input(
         return feat, y, types
 
     def objective(
-        self,
-        trial: Any,  # optuna.trial.Trial
-        dtrain: Any,  # xgboost.DMatrix
-        dvalid: Any,  # xgboost.DMatrix
-        num_boost_round: int,
-        device: str = 'cpu'
-    ) -> float:
+            self,
+            trial: Any,  # optuna.trial.Trial
+            dtrain: Any,  # xgboost.DMatrix
+            dvalid: Any,  # xgboost.DMatrix
+            num_boost_round: int,
+            device: str = 'cpu') -> float:
         r"""Objective function to be optimized.
 
         Args:
@@ -123,7 +122,8 @@ def objective(
             "alpha":
             (0.0 if not trial.suggest_categorical('use_alpha', [True, False])
              else trial.suggest_float('alpha', 1e-8, 1e2, log=True)),
-            "device": device
+            "device":
+            device
         }
         if device.startswith("gpu") or device.startswith("cuda"):
             self.params["tree_method"] = "hist"
@@ -190,14 +190,9 @@ def objective(
                                     torch.from_numpy(pred))
         return score
 
-    def _tune(
-        self,
-        tf_train: TensorFrame,
-        tf_val: TensorFrame,
-        num_trials: int,
-        num_boost_round: int = 2000,
-        device: str = 'cpu'
-    ):
+    def _tune(self, tf_train: TensorFrame, tf_val: TensorFrame,
+              num_trials: int, num_boost_round: int = 2000,
+              device: str = 'cpu'):
         import optuna
         import xgboost