code: def with_strata(shelter):
      baseline_cols = [
          "(" + c1 + ", " + c2 + ")"
          for c1 in ["CONFISCATED", "OWNER SURRENDER", "STRAY"]
          for c2 in [
              "NONE",
              "TREATABLE MANAGEABLE",
              "TREATABLE REHABILITABLE",
              "UNHEALTHY UNTREATABLE",
          ]
      ]
  
      check_cox(
          shelter,
          x=[
              "intake_condition1",
              "intake_condition2",
              "intake_type",
              "animal_breed",
              "chip_status",
              "surv_hours",
          ],
          expected_hazard_cols=baseline_cols,
          expected_survival_cols=baseline_cols,
          stratify_by=["intake_type", "intake_condition1"],
      )
  
file_path: h2o-3/h2o-py/tests/testdir_algos/coxph/pyunit_coxph_baseline_names_shelter.py
imports: import sys
  import h2o
  from lifelines import CoxPHFitter
  from pandas.testing import assert_frame_equal
  from pandas import read_csv
  import numpy
  from tests import pyunit_utils
  from h2o.estimators.coxph import H2OCoxProportionalHazardsEstimator
label_desc: Use image node_id 3 for calling a global function with example usage: with_strata(shelter) without return types