code: def add_chosen_arm_features(
      batch: Union[CBInput, List[CBInput]]
  ) -> Union[CBInput, List[CBInput]]:
      """
      Add the features for chosen arms to the batch.
      For joint models:
          - Both input and output are`CBInput` objects
          - batch.all_arm_features is a 3D tensor of shape (batch_size, num_arms, arm_dim)
          - `batch.action` is used to choose which features to extract.
      For joint models:
           - Both input and output are `List[CBInput]` objects of len `num_arms`
           - batch.all_arm_features is a 2D tensor of shape (batch_size, arm_dim)
           - This function just extracts the `all_arm_features` attribute from batch and packages them back into a list
  
      Args:
          batch: A batch of input data.
              Attributes of the batch:
                  all_arm_features: Tensor with features of all available arms.
                  action: 2D Tensor of shape (batch_size, 1) with dtype long. For each observation
                      it holds the index of the chosen arm.
      Returns:
          For joint models:
              A 2D Tensor of shape (batch_size, arm_dim) with features of chosen arms.
          For disjoint models:
              A list of 2D Tensors of shape (batch_size, arm_dim)
      """
      if isinstance(batch, CBInput):
          assert batch.context_arm_features.ndim == 3
          assert batch.action is not None
          batch = replace(
              batch,
              features_of_chosen_arm=torch.gather(
                  batch.context_arm_features,
                  1,
                  batch.action.unsqueeze(-1).expand(
                      -1, 1, batch.context_arm_features.shape[2]
                  ),
              ).squeeze(1),
          )
          if batch.arms is not None:
              assert batch.action is not None
              batch = replace(
                  batch,
                  chosen_arm_id=torch.gather(
                      batch.arms, 1, batch.action
                  ),
              )
          return batch
      elif isinstance(batch, list):
          assert isinstance(batch[0], CBInput)
          assert batch[0].context_arm_features.ndim == 2
          return [
              replace(b, features_of_chosen_arm=b.context_arm_features)
              for b in batch
          ]
      else:
          raise ValueError(
              f"Unexpected input type {type(batch)} for _add_chosen_arm_features"
          )
  
file_path: ReAgent/reagent/training/cb/utils.py
imports: import logging
  from dataclasses import replace
  from typing import List, Optional, Union
  import torch
  from reagent.core.types import CBInput
label_desc: Use image node_id 1 for calling a global function with example usage: add_chosen_arm_features(batch) and returns: batch, unknown