def remove_batch_from_layout(layout):
    """
    The tf-mesh layout splits across batch size, remove it.
    Useful for prediction steps, when you no longer want large batches.

    :param layout: string describing tf-mesh layout
    :return: layout minus batch dimension
    """
    layout = layout.split(',')
    ret_layout = ""
    for i in layout:
        if "batch" in i:
            pass
        else:
            ret_layout += f"{i},"
    return ret_layout[:-1]
The tf-mesh layout splits across batch size, remove it.
Useful for prediction steps, when you no longer want large batches.

:param layout: string describing tf-mesh layout
:return: layout minus batch dimension
utils.py