xora/utils/dist_util.py

from enum import Enum

class AccelerationType(Enum):
    CPU = "cpu"
    GPU = "gpu"
    TPU = "tpu"
    MPS = "mps"

def execute_graph() -> None:
    if _acceleration_type == AccelerationType.TPU:
        xm.mark_step()