def get_time_axis(layout):
    """
    Extract the time axis from the layout, for example for breaking sequence into
    segments.
    """
    if layout in ["TB", "TBD"]:
        return 0
    if layout in ["BT", "BTD"]:
        return 1
    if layout in ["BCTD"]:
        return 2
    raise ValueError("Unsupported layout = {}".format(layout))