code: def _sub_policy_to_probability_map(policy: Policy) -> _DataNode:
      sub_policies = policy.sub_policies
      max_policy_len = max(
          len(sub_policy) for sub_policy in sub_policies
      )
      prob = np.array(
          [
              [0.0 for _ in range(max_policy_len)]
              for _ in range(len(sub_policies))
          ],
          dtype=np.float32,
      )
      for sub_policy_id, sub_policy in enumerate(sub_policies):
          for stage_idx, (aug_name, p, mag) in enumerate(sub_policy):
              prob[sub_policy_id, stage_idx] = p
      return types.Constant(prob)
  
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 8 for calling a global function with example usage: _sub_policy_to_probability_map(policy) and returns: types