code: def test_onehot_category_class_count(total_classes: int):
      # See this for reasoning: https://github.com/rapidsai/cuml/issues/2690
      # All tests use sparse=True to avoid memory errors
  
      encoder = OneHotEncoder(handle_unknown="ignore", sparse=True)
  
      # ==== 2 Rows ====
      example_df = DataFrame()
      example_df["high_cardinality_column"] = cp.linspace(
          0, total_classes - 1, total_classes
      )
      example_df["low_cardinality_column"] = ["A"] * 200 + ["B"] * (
          total_classes - 200
      )
  
      assert (
          encoder.fit_transform(example_df).shape[1]
          == total_classes + 2
      )
  
      # ==== 3 Rows ====
      example_df = DataFrame()
      example_df["high_cardinality_column"] = cp.linspace(
          0, total_classes - 1, total_classes
      )
      example_df["low_cardinality_column"] = ["A"] * total_classes
      example_df["med_cardinality_column"] = ["B"] * total_classes
  
      assert (
          encoder.fit_transform(example_df).shape[1]
          == total_classes + 2
      )
  
      # ==== N Rows (Even Split) ====
      num_rows = [3, 10, 100]
  
      for row_count in num_rows:
          class_per_row = (
              int(math.ceil(total_classes / float(row_count))) + 1
          )
          example_df = DataFrame()
  
          for row_idx in range(row_count):
              example_df[str(row_idx)] = cp.linspace(
                  row_idx * class_per_row,
                  ((row_idx + 1) * class_per_row) - 1,
                  class_per_row,
              )
  
          assert (
              encoder.fit_transform(example_df).shape[1]
              == class_per_row * row_count
          )
  
imports: from sklearn.preprocessing import OneHotEncoder
  from cuml.testing.utils import stress_param, from_df_to_numpy, assert_inverse_equal, generate_inputs_from_categories
  from cuml.preprocessing import OneHotEncoder
  from cuml.internals.safe_imports import gpu_only_import_from
  import pytest
  from cuml.internals.safe_imports import cpu_only_import
  import math
  from cuml.internals.safe_imports import gpu_only_import
label_desc: Use image node_id 17 for calling a global function with example usage: test_onehot_category_class_count(total_classes) without return types