xma.accelerator

class Accelerator(*values)

Bases: Enum

cpu = 'cpu'

cuda = 'cuda'

static get_accelerator() → Accelerator

static get_current_device() → int | str

static get_kernel_backend() → KernelBackend

static get_sm_count(device: device | None = None) → int

rocm = 'rocm'

static synchronize() → None

tpu = 'tpu'

trainium = 'trainium'

class KernelBackend(*values)

Bases: Enum

cuda = 'cuda'

get_compatible_accelerator() → Accelerator

nki = 'nki'

pallas = 'pallas'

rocm = 'rocm'

torch = 'torch'

triton = 'triton'

verify_accelerator() → bool