code: def create_discrete_action_masking_layer(all_logits, action_masks, action_size):
          
          action_idx = [0] + list(np.cumsum(action_size))
          branches_logits = [all_logits[:, action_idx[i]:action_idx[i + 1]] for i in range(len(action_size))]
          branch_masks = [action_masks[:, action_idx[i]:action_idx[i + 1]] for i in range(len(action_size))]
          raw_probs = [tf.multiply(tf.nn.softmax(branches_logits[k]) + 1.0e-10, branch_masks[k])
                       for k in range(len(action_size))]
          normalized_probs = [
              tf.divide(raw_probs[k], tf.reduce_sum(raw_probs[k], axis=1, keepdims=True))
              for k in range(len(action_size))]
          output = tf.concat([tf.multinomial(tf.log(normalized_probs[k]), 1) for k in range(len(action_size))], axis=1)
          return output, tf.concat([tf.log(normalized_probs[k] + 1.0e-10) for k in range(len(action_size))], axis=1)
docs: Creates a masking layer for the discrete actions
          :param all_logits: The concatenated unnormalized action probabilities for all branches
          :param action_masks: The mask for the logits. Must be of dimension [None x total_number_of_action]
          :param action_size: A list containing the number of possible actions for each branch
          :return: The action output dimension [batch_size, num_branches] and the concatenated normalized logits