def get_extended_indices(decode_steps, output_vocab_size, model_config):
  """Convert DecodeSteps into a tensor of extended action ids."""
  # This initial value will be broadcast to the length of decode_steps.
  extended_action_indices = tf.constant(0, dtype=tf.int64)
  for action_type_range in _get_action_types_to_range(output_vocab_size,
                                                      model_config):
    is_type = tf.equal(
        tf.constant(action_type_range.action_type, dtype=tf.int64),
        decode_steps.action_types)
    # For each timestep, exactly one of the action_type_ranges will be added,
    # so this sum will populate each entry on exactly one iteration.
    extended_action_indices += (
        tf.to_int64(is_type) *
        (decode_steps.action_ids + action_type_range.start_index))
  return extended_action_indices
Convert DecodeSteps into a tensor of extended action ids.
language/xsp/model/decode_utils.py