code: def _sub_policy_to_augmentation_matrix_map(
      policy: Policy,
  ) -> Tuple[np.ndarray, List[List[_Augmentation]]]:
      """
      Creates a matrix of operators to be called for given sub policy at given stage.
      The output is a tuple `(m, augments)`, where `augments` is a list of augmentations per stage
      - each entry contains a reduced list of unique augmentations used in a corresponding stage.
      The `m` matrix contains the mapping from the original sub_policy_id, to the index within the
      reduced list, for every stage. I.e., for policy `sub_policy_idx`, as the `stage_idx`-ith
      operation in a sequence, the `augments[stage_idx][m[sub_policy_idx][stage_idx]]` operator
      should be called.
      """
      sub_policies = policy.sub_policies
      max_policy_len = max(
          len(sub_policy) for sub_policy in sub_policies
      )
      augmentations = []  # list of augmentations in each stage
      for stage_idx in range(max_policy_len):
          stage_augments = set()
          stage_augments_list = []
          for sub_policy in sub_policies:
              if stage_idx < len(sub_policy):
                  aug, _, _ = sub_policy[stage_idx]
                  if aug not in stage_augments:
                      stage_augments.add(aug)
                      stage_augments_list.append(aug)
          augmentations.append(stage_augments_list + [a.identity])
      identity_id = [
          len(stage_augments) - 1 for stage_augments in augmentations
      ]
      augment_to_id = [
          {
              augmentation: i
              for i, augmentation in enumerate(stage_augments)
          }
          for stage_augments in augmentations
      ]
      augments_by_id = np.array(
          [
              [
                  identity_id[stage_idx]
                  for stage_idx in range(max_policy_len)
              ]
              for _ in range(len(sub_policies))
          ],
          dtype=np.int32,
      )
      for sub_policy_id, sub_policy in enumerate(sub_policies):
          for stage_idx, (augment, p, mag) in enumerate(sub_policy):
              augments_by_id[sub_policy_id, stage_idx] = augment_to_id[
                  stage_idx
              ][augment]
      return augments_by_id, augmentations
  
file_path: DALI/dali/python/nvidia/dali/auto_aug/auto_augment.py
imports: from typing import List, Optional, Tuple, Union
  from nvidia.dali import fn
  from nvidia.dali import types
  from nvidia.dali.auto_aug import augmentations
  from nvidia.dali.auto_aug.core import _Augmentation, Policy, signed_bin
  from nvidia.dali.auto_aug.core._args import forbid_unused_kwargs
  from nvidia.dali.auto_aug.core._utils import get_translations, pretty_select
  from nvidia.dali.data_node import DataNode
label_desc: Use image node_id 10 for calling a global function with example usage: _sub_policy_to_augmentation_matrix_map(policy) and returns: augments_by_id,  augmentations