code: def test_sub_policy(randomly_negate, dev, batch_size):
      num_magnitude_bins = 10
  
      @augmentation(
          mag_range=(0, 9),
          mag_to_param=mag_to_param_with_op_id(1),
          param_device=dev,
      )
      def first(data, op_id_mag_id):
          return fn.cat(data, op_id_mag_id)
  
      @augmentation(
          mag_range=(10, 19),
          mag_to_param=mag_to_param_with_op_id(2),
          randomly_negate=randomly_negate,
          param_device=dev,
      )
      def second(data, op_id_mag_id):
          return fn.cat(data, op_id_mag_id)
  
      @augmentation(
          mag_range=(20, 29),
          mag_to_param=mag_to_param_with_op_id(3),
          randomly_negate=randomly_negate,
          param_device=dev,
      )
      def third(data, op_id_mag_id):
          return fn.cat(data, op_id_mag_id)
  
      sub_policies = [
          [(first, 1, 0), (second, 1, 5), (third, 1, 3)],
          [(first, 1, 1), (third, 1, 4), (first, 1, 2)],
          [(second, 1, 2), (first, 1, 3), (third, 1, 4)],
          [(second, 1, 3), (third, 1, 2), (first, 1, 5)],
          [(third, 1, 4), (first, 1, 1), (second, 1, 1)],
          [(third, 1, 5), (second, 1, 9), (first, 1, 2)],
          [(first, 1, 6), (first, 1, 1)],
          [(third, 1, 7)],
          [
              (first, 1, 8),
              (first, 1, 4),
              (second, 1, 7),
              (second, 1, 6),
          ],
      ]
  
      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()
  
      sub_policy_outputs = collect_sub_policy_outputs(
          sub_policies, num_magnitude_bins
      )
      # magnitudes are chosen so that the magnitude of the first op in
      # each sub-policy identifies the sub-policy
      assert len({out[0][1] for out in sub_policy_outputs}) == len(
          sub_policy_outputs
      )
      output_cases = {
          out[0][1]: np.array(out) for out in sub_policy_outputs
      }
  
      sub_policy_negation_cases = []
      for sub_policy in sub_policies:
          negated = []
          for aug, _, _ in sub_policy:
              if aug.randomly_negate:
                  negated.append((True, False))
              else:
                  negated.append((False,))
          sub_policy_negation_cases.append(
              list(itertools.product(*negated))
          )
      assert len(sub_policy_outputs) == len(sub_policy_negation_cases)
  
      for _ in range(5):
          (output,) = p.run()
          if dev == "gpu":
              output = output.as_cpu()
          output = [np.array(sample) for sample in output]
          for sample in output:
              test_sample = (
                  sample if not randomly_negate else np.abs(sample)
              )
              np.testing.assert_equal(
                  np.abs(test_sample), output_cases[test_sample[0][1]]
              )
              for op_mag in sample:
                  if op_mag[1] < 0:
                      # the `second` and `third` augmentation are marked as randomly_negated
                      assert op_mag[0] in [2, 3], f"{sample}"
          if randomly_negate:
              # for each sub-policy, count occurrences of any possible sequence
              # of magnitude signs
              negation_cases = {
                  out[0][1]: {case: 0 for case in cases}
                  for out, cases in zip(
                      sub_policy_outputs, sub_policy_negation_cases
                  )
              }
              for sample in output:
                  mag_signs = tuple(op_mag[1] < 0 for op_mag in sample)
                  negation_cases[np.abs(sample[0][1])][mag_signs] += 1
              counts, expected_counts = [], []
              for sub_policy_cases in negation_cases.values():
                  expected = batch_size / (
                      len(sub_policies) * len(sub_policy_cases)
                  )
                  for count in sub_policy_cases.values():
                      counts.append(count)
                      expected_counts.append(expected)
              stat = chisquare(counts, expected_counts)
              # 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 5 for calling a global function with example usage: test_sub_policy(randomly_negate, dev, batch_size) and returns: fn, fn, fn