code: def split_ohe(
      X_ohe: "Union[np.ndarray, torch.Tensor, tf.Tensor]",
      category_map: Dict[int, List[str]],
  ) -> Tuple[List, List]:
      """
      Splits a one-hot encoding array in a list of numerical heads and a list of categorical heads. Since by
      convention the numerical heads are merged in a single head, if the function returns a list of numerical heads,
      then the size of the list is 1.
  
      Parameters
      ----------
      X_ohe
          One-hot encoding representation. This can be any type of tensor: `np.ndarray`, `torch.Tensor`, `tf.Tensor`.
      category_map
          Dictionary of category mapping. The keys are column indexes and the values are lists containing the
          possible values of a feature.
  
      Returns
      -------
      X_ohe_num_split
          List of numerical heads. If different than ``None``, the list's size is 1.
      X_ohe_cat_split
          List of categorical one-hot encoded heads.
      """
      assert hasattr(
          X_ohe, "shape"
      ), "X_ohe needs to have `shape` attribute."
      X_ohe_num_split, X_ohe_cat_split = [], []
      offset = 0
  
      # Compute the number of columns spanned by the categorical one-hot encoded heads, and the number of columns
      # spanned by the numerical heads.
      cat_feat = int(
          np.sum([len(vals) for vals in category_map.values()])
      )
      num_feat = X_ohe.shape[1] - cat_feat
  
      # If the number of numerical features is different than 0, then the first `num_feat` columns correspond
      # to the numerical features
      if num_feat > 0:
          X_ohe_num_split.append(X_ohe[:, :num_feat])
          offset = num_feat
  
      # If there exist categorical features, then extract them one by one
      if cat_feat > 0:
          for id in sorted(category_map.keys()):
              X_ohe_cat_split.append(
                  X_ohe[:, offset : offset + len(category_map[id])]
              )
              offset += len(category_map[id])
  
      return X_ohe_num_split, X_ohe_cat_split
  
imports: from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple, Union
  import numpy
  import pandas
  from scipy.special import softmax
  from sklearn.compose import ColumnTransformer
  from sklearn.preprocessing import OneHotEncoder, StandardScaler
label_desc: Use image node_id 2 for calling a global function with example usage: split_ohe(X_ohe, category_map) and returns: X_ohe_num_split,  X_ohe_cat_split