code: def test_regressor_sg_train_mg_predict(
      datatype, keys, data_size, fit_intercept, client
  ):
      # Just testing for basic compatibility w/ dask-ml's ParallelPostFit.
      # Refer to test_pickle.py for more extensive testing of single-GPU
      # model serialization.
  
      nrows, ncols, n_info = data_size
      X_train, y_train, _ = make_dataset(datatype, nrows, ncols, n_info)
  
      X_train_local = X_train.compute()
      y_train_local = y_train.compute()
  
      local_model = cuml.linear_model.LinearRegression(
          fit_intercept=fit_intercept
      )
      local_model.fit(X_train_local, y_train_local)
  
      dist_model = ParallelPostFit(estimator=local_model)
  
      predictions = dist_model.predict(X_train).compute()
  
      assert isinstance(predictions, cupy.ndarray)
  
      # Dataset should be fairly linear already so the predictions should
      # be very close to the training data.
      np.testing.assert_allclose(
          predictions.get(),
          y_train.compute().get(),
          atol=1e-3,
          rtol=1e-3,
      )
  
file_path: cuml/python/cuml/tests/dask/test_dask_base.py
imports: from sklearn.model_selection import train_test_split
  from cuml.dask.datasets import make_regression
  from cuml.dask.linear_model import LinearRegression
  from cuml.internals.safe_imports import cpu_only_import_from
  import pytest
  from cuml.internals.safe_imports import cpu_only_import
  import cuml
  from cuml.dask.datasets import make_blobs
  from cuml.testing.dask.utils import load_text_corpus
  from cuml.dask.naive_bayes.naive_bayes import MultinomialNB
  from cuml.dask.cluster import KMeans
  from dask_ml.wrappers import ParallelPostFit
  from cuml.internals.safe_imports import gpu_only_import
label_desc: Use image node_id 5 for calling a global function with example usage: test_regressor_sg_train_mg_predict(datatype, keys, data_size, fit_intercept, client) without return types