import torch


def _get_chosen_action_features(
    all_action_features: torch.Tensor, chosen_actions: torch.Tensor
) -> torch.Tensor:
    """
    Pick the features for chosen actions out of a tensor with features of all actions

    Args:
        all_action_features: 3D Tensor of shape (batch_size, num_actions, action_dim) with
            features of all available actions.
        chosen_actions: 2D Tensor of shape (batch_size, 1) with dtype long. For each observation
            it holds the index of the chosen action.
    Returns:
        A 2D Tensor of shape (batch_size, action_dim) with features of chosen actions.
    """
    assert all_action_features.ndim == 3
    return torch.gather(
        all_action_features,
        1,
        chosen_actions.unsqueeze(-1).expand(-1, 1, all_action_features.shape[2]),
    ).squeeze(1)