code: def test_divide_by_mask(batch_size, n_obs, prop_true, dtype):
      """Test the helper that splits a dataset in 2 based on a boolean mask"""
      # Generate random data, mask and batch indices
      data_np = (
          (np.random.uniform(-1.0, 1.0, (batch_size, n_obs)))
          .astype(dtype)
          .transpose()
      )
      nb_true = int(prop_true * batch_size)
      mask_np = np.random.permutation(
          [False] * (batch_size - nb_true) + [True] * nb_true
      )
      b_id_np = np.array(range(batch_size), dtype=np.int32)
      data, *_ = input_to_cuml_array(data_np)
      mask, *_ = input_to_cuml_array(mask_np)
      b_id, *_ = input_to_cuml_array(b_id_np)
  
      # Call the tested function
      sub_data, sub_id = [None, None], [None, None]
      (
          sub_data[0],
          sub_id[0],
          sub_data[1],
          sub_id[1],
      ) = auto_arima._divide_by_mask(data, mask, b_id)
  
      # Compute the expected results in pure Python
      sub_data_ref = [
          data_np[:, np.logical_not(mask_np)],
          data_np[:, mask_np],
      ]
      sub_id_ref = [b_id_np[np.logical_not(mask_np)], b_id_np[mask_np]]
  
      # Compare the results
      for i in range(2):
          # First check the cases of empty sub-batches
          if sub_data[i] is None:
              # The reference must be empty
              assert sub_data_ref[i].shape[1] == 0
              # And the id array must be None too
              assert sub_id[i] is None
          # When the sub-batch is not empty, compare to the reference
          else:
              np.testing.assert_allclose(
                  sub_data[i].to_output("numpy"), sub_data_ref[i]
              )
              np.testing.assert_array_equal(
                  sub_id[i].to_output("numpy"), sub_id_ref[i]
              )
  
file_path: cuml/python/cuml/tests/test_auto_arima.py
imports: from cuml.internals.input_utils import input_to_cuml_array
  from cuml.tsa import auto_arima
  import pytest
  from cuml.internals.safe_imports import cpu_only_import
label_desc: Use image node_id 2 for calling a global function with example usage: test_divide_by_mask(batch_size, n_obs, prop_true, dtype) without return types