code: def test_n_classes_small(client):
      def assert_small(X, y, n_classes):
          X_df, y_df = _prep_training_data(
              client, X, y, partitions_per_worker=1
          )
          from cuml.dask.linear_model import (
              LogisticRegression as cumlLBFGS_dask,
          )
  
          lr = cumlLBFGS_dask()
          lr.fit(X_df, y_df)
          assert lr._num_classes == n_classes
          return lr
  
      X = np.array([(1, 2), (1, 3)], np.float32)
      y = np.array([1.0, 0.0], np.float32)
      lr = assert_small(X=X, y=y, n_classes=2)
      assert np.array_equal(
          lr.classes_.to_numpy(), np.array([0.0, 1.0], np.float32)
      )
  
      X = np.array([(1, 2), (1, 3), (1, 2.5)], np.float32)
      y = np.array([1.0, 0.0, 1.0], np.float32)
      lr = assert_small(X=X, y=y, n_classes=2)
      assert np.array_equal(
          lr.classes_.to_numpy(), np.array([0.0, 1.0], np.float32)
      )
  
      X = np.array([(1, 2), (1, 2.5), (1, 3)], np.float32)
      y = np.array([1.0, 1.0, 0.0], np.float32)
      lr = assert_small(X=X, y=y, n_classes=2)
      assert np.array_equal(
          lr.classes_.to_numpy(), np.array([0.0, 1.0], np.float32)
      )
  
      X = np.array([(1, 2), (1, 3), (1, 2.5)], np.float32)
      y = np.array([10.0, 50.0, 20.0], np.float32)
      lr = assert_small(X=X, y=y, n_classes=3)
      assert np.array_equal(
          lr.classes_.to_numpy(),
          np.array([10.0, 20.0, 50.0], np.float32),
      )
  
imports: from cuml.internals.safe_imports import gpu_only_import
  import pytest
  from cuml.dask.common import utils
  from functools import partial
  from sklearn.metrics import accuracy_score, mean_squared_error
  from sklearn.datasets import make_classification
  from sklearn.linear_model import LogisticRegression
  from cuml.internals.safe_imports import cpu_only_import
  from cuml.testing.utils import array_equal
  from scipy.sparse import csr_matrix
label_desc: Use image node_id 10 for calling a global function with example usage: test_n_classes_small(client) and returns: lr