code: def test_split_ohe(dataset):
      """Test the one-hot encoding splitting of a dataset."""
  
      # Unpack data
      X = dataset["X_train"]
      feature_names = dataset["metadata"]["feature_names"]
      category_map = dataset["metadata"].get("category_map", {})
  
      # Get heterogeneous preprocessor
      preprocessor, _ = get_he_preprocessor(
          X=X, feature_names=feature_names, category_map=category_map
      )
  
      # Preprocess the dataset
      X_ohe = preprocessor(X)
  
      # Split the ohe representation in multiple heads
      num_heads, cat_heads = split_ohe(
          X_ohe=X_ohe, category_map=category_map
      )
  
      # Check the number of numerical heads.
      if len(category_map) < len(feature_names):
          assert len(num_heads) == 1
      else:
          assert len(num_heads) == 0
  
      # Check the number of categorical heads.
      assert len(cat_heads) == len(category_map)
  
imports: from typing import Union, List
  import pytest
  from pytest_lazyfixture import lazy_fixture
  import numpy
  from numpy.testing import assert_allclose
  import tensorflow
  import tensorflow.keras
  from alibi.explainers import CounterfactualRLTabular
  from alibi.explainers.backends.cfrl_base import get_hard_distribution
  from alibi.explainers.backends.cfrl_tabular import get_he_preprocessor, split_ohe, get_numerical_conditional_vector, get_categorical_conditional_vector, get_statistics, get_conditional_vector, sample
label_desc: Use image node_id 2 for calling a global function with example usage: test_split_ohe(dataset) without return types