code: def test_op_skipping(dev):
      num_magnitude_bins = 20
      batch_size = 2400
  
      @augmentation(
          mag_range=(0, 19),
          mag_to_param=mag_to_param_with_op_id(1),
          randomly_negate=True,
          param_device=dev,
      )
      def first(data, op_id_mag_id):
          return fn.cat(data, op_id_mag_id)
  
      @augmentation(
          mag_range=(0, 19),
          mag_to_param=mag_to_param_with_op_id(2),
          randomly_negate=True,
          param_device=dev,
      )
      def second(data, op_id_mag_id):
          return fn.cat(data, op_id_mag_id)
  
      @augmentation(
          mag_range=(0, 19),
          mag_to_param=mag_to_param_with_op_id(3),
          param_device=dev,
      )
      def third(data, op_id_mag_id):
          return fn.cat(data, op_id_mag_id)
  
      @augmentation(
          mag_range=(0, 19),
          mag_to_param=mag_to_param_with_op_id(4),
          param_device=dev,
      )
      def first_stage_only(data, op_id_mag_id):
          return fn.cat(data, op_id_mag_id)
  
      @augmentation(
          mag_range=(0, 19),
          mag_to_param=mag_to_param_with_op_id(5),
          param_device=dev,
      )
      def second_stage_only(data, op_id_mag_id):
          return fn.cat(data, op_id_mag_id)
  
      sub_policies = [
          [(first, 0.5, 1), (first, 0.25, 2)],
          [(second, 0.8, 3), (second, 0.7, 4)],
          [(first, 0.9, 5), (second, 0.6, 6)],
          [(second, 0.3, 7), (first, 0.25, 8)],
          [(third, 1, 9), (third, 0.75, 10)],
          [(third, 0.3, 11), (first, 0.22, 12)],
          [(second, 0.6, 13), (third, 0, 14)],
          [(first_stage_only, 0.5, 15), (third, 0.7, 16)],
          [(third, 0.8, 17), (second_stage_only, 0.6, 18)],
      ]
  
      # sub_policy_cases = [[] for _ in range(len(sub_policies))]
      expected_counts = {tuple(): 0.0}
      for (left_aug, left_p, left_mag), (
          right_aug,
          right_p,
          right_mag,
      ) in sub_policies:
          expected_counts[tuple()] += (
              (1.0 - left_p) * (1 - right_p) / len(sub_policies)
          )
          only_left_p = left_p * (1 - right_p) / len(sub_policies)
          only_right_p = (1 - left_p) * right_p / len(sub_policies)
          for aug, mag, prob in [
              (left_aug, left_mag, only_left_p),
              (right_aug, right_mag, only_right_p),
          ]:
              if not aug.randomly_negate:
                  expected_counts[(mag,)] = prob
              else:
                  expected_counts[(mag,)] = prob / 2
                  expected_counts[(-mag,)] = prob / 2
          sign_cases = [
              (-1, 1) if aug.randomly_negate else (1,)
              for aug in (left_aug, right_aug)
          ]
          sign_cases = list(itertools.product(*sign_cases))
          prob = left_p * right_p / len(sub_policies)
          for left_sign, right_sign in sign_cases:
              mags = (left_sign * left_mag, right_sign * right_mag)
              expected_counts[mags] = prob / len(sign_cases)
      expected_counts = {
          mag: prob * batch_size
          for mag, prob in expected_counts.items()
          if prob > 0
      }
      assert all(
          num_elements >= 5 for num_elements in expected_counts.values()
      ), (
          f"The batch size is too small (i.e. some output cases are expected less "
          f"than five times in the output): {expected_counts}"
      )
  
      policy = Policy(
          "MyPolicy",
          num_magnitude_bins=num_magnitude_bins,
          sub_policies=sub_policies,
      )
      p = concat_aug_pipeline(
          batch_size=batch_size, dev=dev, policy=policy
      )
      p.build()
  
      for _ in range(5):
          (output,) = p.run()
          if dev == "gpu":
              output = output.as_cpu()
          output = [np.array(sample) for sample in output]
          actual_counts = {
              allowed_case: 0 for allowed_case in expected_counts
          }
          for sample in output:
              mags = tuple(int(op_mag[1]) for op_mag in sample)
              actual_counts[mags] += 1
  
          actual, expected = [], []
          for mags in expected_counts:
              actual.append(actual_counts[mags])
              expected.append(expected_counts[mags])
          stat = chisquare(actual, expected)
          # assert that the magnitudes negation looks independently enough
          # (0.05 <=), but also that it is not too ideal (i.e. like all
          # cases happening exactly the expected number of times)
          assert 0.05 <= stat.pvalue <= 0.95, f"{stat}"
  
imports: import itertools
  import os
  import unittest
  import random
  import numpy
  from scipy.stats import chisquare
  from nose2.tools import params
  from nvidia.dali import fn, types
  from nvidia.dali import pipeline_def
  from nvidia.dali.auto_aug import auto_augment, augmentations
  from nvidia.dali.auto_aug.core import augmentation, Policy
  from test_utils import get_dali_extra_path, check_batch
  from nose_utils import assert_raises, assert_warns
label_desc: Use image node_id 6 for calling a global function with example usage: test_op_skipping(dev) and returns: fn, fn, fn, fn, fn