code: def test_ops_selection_and_mags(case_idx, args):
      dev, use_sign, n, num_ops = args
      num_magnitude_bins = 9
      # the chisquare expects at least 5 elements in a bin and we can have around
      # (num_ops * (2**use_signs)) ** n ops
      batch_size = 2048
      magnitude_cases = list(range(num_magnitude_bins))
      m = magnitude_cases[case_idx % len(magnitude_cases)]
  
      def mag_to_param_with_op_id(op_id):
          def mag_to_param(magnitude):
              return np.array([op_id, magnitude], dtype=np.int32)
  
          return mag_to_param
  
      @augmentation(param_device=dev)
      def op(data, op_id_mag_id):
          return fn.cat(data, op_id_mag_id)
  
      augmentations = [
          op.augmentation(
              mag_range=(10 * i + 1, 10 * i + num_magnitude_bins),
              mag_to_param=mag_to_param_with_op_id(i + 1),
              randomly_negate=use_sign and i % 3 == 0,
          )
          for i in range(num_ops)
      ]
  
      expected_counts = {}
      seq_prob = 1.0 / (num_ops**n)
      for aug_sequence in itertools.product(*([augmentations] * n)):
          possible_signs = [
              (-1, 1) if aug.randomly_negate else (1,)
              for aug in aug_sequence
          ]
          possible_signs = tuple(itertools.product(*possible_signs))
          prob = seq_prob / len(possible_signs)
          for signs in possible_signs:
              assert len(aug_sequence) == len(signs)
              outs = []
              for aug, sign in zip(aug_sequence, signs):
                  mag = aug._get_magnitudes(num_magnitude_bins)[m]
                  op_id_mag = aug.mag_to_param(mag * sign)
                  outs.append(op_id_mag)
              expected_counts[
                  tuple(el for out in outs for el in out)
              ] = prob
      expected_counts = {
          output: p * batch_size
          for output, p in expected_counts.items()
      }
  
      @pipeline_def(
          enable_conditionals=True,
          batch_size=batch_size,
          num_threads=4,
          device_id=0,
          seed=42,
      )
      def pipeline():
          data = types.Constant([], dtype=types.INT32)
          if dev == "gpu":
              data = data.gpu()
          data = rand_augment.apply_rand_augment(
              augmentations,
              data,
              n=n,
              m=m,
              num_magnitude_bins=num_magnitude_bins,
          )
          return fn.reshape(data, shape=(-1, 2))
  
      p = pipeline()
      p.build()
      for i in range(3):
          (output,) = p.run()
          output = [
              np.array(s)
              for s in (output.as_cpu() if dev == "gpu" else output)
          ]
          actual_count = {
              allowed_out: 0 for allowed_out in expected_counts
          }
          for sample in output:
              assert len(sample) == n, f"{i} {sample}"
              out = tuple(el for op_mag in sample for el in op_mag)
              actual_count[out] += 1
          actual = []
          expected = []
          for out in expected_counts:
              actual.append(actual_count[out])
              expected.append(expected_counts[out])
          stat = chisquare(actual, expected)
          assert (
              0.01 <= stat.pvalue <= 0.99
          ), f"{stat} {actual} {expected}"
  
imports: import itertools
  import os
  import random
  import unittest
  import numpy
  from scipy.stats import chisquare
  from nose2.tools import params
  from nvidia.dali import fn, tensors, types
  from nvidia.dali import pipeline_def
  from nvidia.dali.auto_aug import rand_augment
  from nvidia.dali.auto_aug.core import augmentation
  from test_utils import get_dali_extra_path, check_batch
  from nose_utils import assert_raises
label_desc: Use image node_id 3 for calling a global function with example usage: test_ops_selection_and_mags(case_idx, args) and returns: mag_to_param, fn, fn, np