import torch


def filter_by_for_lists(action_fields, target_list):
    """Returns the subset of `target` corresponding to `action_fields`.

    Autoregressive actions are composed of many logits.  We often want to select a
    subset of these logits.

    Args:
      action_fields: One of 'action_type', 'delay', or 'arguments'.
      target: A list of tensors corresponding to the SC2 action spec. [T x B x S]
    Returns:
      A tensor corresponding to a subset of `target`, with only the tensors relevant
      to `action_fields`.
    """
    return torch.cat([getattr(b, action_fields) for a in target_list for b in a], dim=0)