def get_batch_axis(layout):
    """
    Extract the batch axis from the layout
    """
    if layout in ["TB", "TBD"]:
        return 1
    if layout in ["BT", "BTD", "BCTD"]:
        return 0
    raise ValueError("Unsupported layout = {}".format(layout))

    Extract the batch axis from the layout
    
understand_icl/fairseq/fairseq/models/speech_to_text/utils.py