from typing import Union
from typing import List
from typing import Any
from typing import Dict


def get_sample_from_batched_tensors(tensors: Union[List[Any], Dict[str, Any]],
                                    idx: int) -> Dict[str, Any]:
  """Collect sample at index idx of a given dictionary or list.

  Args:
    tensors: Rowified or columnarized format of a batch of inputs.
    idx: Index to fetch the batch from
  Returns:
    Same dictionary where idx is chosen from values.
  """
  if isinstance(tensors, list):
    return tensors[idx] if idx < len(tensors) else {}
  return {key: val[idx] for key, val in tensors.items()}