# Buggy code
def _wrap_train_tuple(data, label, sample_weight=None, init_score=None):
    return TrainTuple(data, label, sample_weight, init_score)


# Traceback (TypeError)
In [1]: from mars.learn.contrib import lightgbm as lgb
/Users/qinxuye/miniconda3/envs/mars3.6/lib/python3.6/site-packages/lightgbm/__init__.py:48: UserWarning: Starting from version 2.2.1, the library file in distribution wheels for macOS is built by the Apple Clang (Xcode_8.3.3) compiler.
This means that in case of installing LightGBM from PyPI via the ``pip install lightgbm`` command, you don't need to install the gcc compiler anymore.
Instead of that, you need to install the OpenMP library, which is required for running LightGBM on the system with the Apple Clang compiler.
You can install the OpenMP library by the following command: ``brew install libomp``.
"You can install the OpenMP library by the following command: ``brew install libomp``.", UserWarning)

In [2]: lg_reg = lgb.LGBMRegressor(colsample_bytree=0.3, learning_rate=0.1,
...:                            max_depth=5, reg_alpha=10, n_estimators=10)

In [3]: from sklearn.datasets import make_classification

In [4]: x, y = make_classification()

In [6]: lg_reg.fit(x, y)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-6-649d7fa4c388> in <module>
----> 1 lg_reg.fit(x, y)

~/Workspace/mars/mars/learn/contrib/lightgbm/regressor.py in fit(self, X, y, sample_weight, init_score, eval_set, eval_sample_weight, eval_init_score, session, run_kwargs, **kwargs)
30                           eval_sets=self._wrap_eval_tuples(eval_set, eval_sample_weight, eval_init_score),
31                           model_type=LGBMModelType.REGRESSOR,
---> 32                           session=session, run_kwargs=run_kwargs, **kwargs)
33
34             self.set_params(**model.get_params())

~/Workspace/mars/mars/learn/contrib/lightgbm/train.py in train(params, train_set, eval_sets, **kwargs)
323     base_port = kwargs.pop('base_port', None)
324
--> 325     aligns = align_data_set(train_set)
326     for eval_set in eval_sets:
327         aligns += align_data_set(eval_set)

~/Workspace/mars/mars/learn/contrib/lightgbm/align.py in align_data_set(dataset)
104
105 def align_data_set(dataset):
--> 106     out_types = get_output_types(dataset.data, dataset.label, dataset.sample_weight, dataset.init_score)
107     op = LGBMAlign(data=dataset.data, label=dataset.label, sample_weight=dataset.sample_weight,
108                    init_score=dataset.init_score, output_types=out_types)

~/Workspace/mars/mars/core.py in get_output_types(unknown_as, *objs)
891                 output_types.append(unknown_as)
892             else:  # pragma: no cover
--> 893                 raise TypeError('Output can only be tensor, dataframe or series')
894     return output_types

TypeError: Output can only be tensor, dataframe or series

# Fixed code
def _wrap_train_tuple(cls, data, label, sample_weight=None, init_score=None):
    data = cls._convert_tileable(data)
    label = cls._convert_tileable(label)
    sample_weight = cls._convert_tileable(sample_weight)
    init_score = cls._convert_tileable(init_score)
    return TrainTuple(data, label, sample_weight, init_score)