code: def test_stratify_split(
      df_engine, nrows, atol, class_probs, ray_cluster_2cpu
  ):
      npartitions = 10
      thresholds = np.cumsum((class_probs * nrows).astype(int))
  
      df = pd.DataFrame(
          np.random.randint(0, 100, size=(nrows, 3)),
          columns=["A", "B", "C"],
      )
  
      def get_category(v):
          if v < thresholds[0]:
              return 0
          if thresholds[0] <= v < thresholds[1]:
              return 1
          return 2
  
      df["category"] = df.index.map(get_category).astype(np.int8)
  
      if isinstance(df_engine, DaskEngine):
          df = df_engine.df_lib.from_pandas(df, npartitions=npartitions)
  
      probs = (0.7, 0.1, 0.2)
      split_params = {
          "type": "stratify",
          "column": "category",
          "probabilities": probs,
      }
      splitter = get_splitter(**split_params)
  
      backend = Mock()
      backend.df_engine = df_engine
      splits = splitter.split(df, backend, random_seed=42)
      assert len(splits) == 3
  
      ratios = class_probs * nrows
      for split, p in zip(splits, probs):
          if isinstance(df_engine, DaskEngine):
              split = split.compute()
          for idx, r in enumerate(ratios):
              actual = np.sum(split["category"] == idx)
              expected = int(r * p)
              assert np.isclose(actual, expected, atol=atol)
  
      # Test determinism
      splits2 = splitter.split(df, backend, random_seed=7)
      for s1, s2 in zip(splits, splits2):
          if isinstance(df_engine, DaskEngine):
              s1 = s1.compute()
              s2 = s2.compute()
          assert not s1.equals(s2)
  
      splits3 = splitter.split(df, backend, random_seed=42)
      for s1, s3 in zip(splits, splits3):
          if isinstance(df_engine, DaskEngine):
              s1 = s1.compute()
              s3 = s3.compute()
          assert s1.equals(s3)
  
file_path: ludwig/tests/ludwig/data/test_split.py
imports: from datetime import datetime, timedelta
  from itertools import combinations
  from random import randrange
  from unittest.mock import Mock
  import numpy
  import pandas
  import pytest
  from ludwig.data.dataframe.pandas import PandasEngine
  from ludwig.data.split import get_splitter
label_desc: Use image node_id 5 for calling a global function with example usage: test_stratify_split(df_engine, nrows, atol, class_probs, ray_cluster_2cpu) and returns: int, int, int