Device

class DeviceType(Enum):
    CPU = "cpu"
    CUDA = "cuda"
    ROCM = "rocm"
    WEBGPU = "webgpu"
    TPU = "tpu"
    NPU = "npu"
    VULKAN = "vulkan"
    OPENCL = "opencl"
    ONEAPI = "oneapi"
    DIRECTML = "directml"
    ONEDNN = "onednn"
    MKLDNN = "mkldnn"
    TENSORRT = "tensorrt"
    QUANTUM = "quantum"
    ARM = "arm"

class Device:
    def __init__(self, device_type: DeviceType, index: int = 0)

@property
def type(self) -> DeviceType

@property
def index(self) -> int

@property
def name(self) -> str

@classmethod
def device_count(cls) -> int

def get_device(device_str: str) -> Device

def get_device_count(device_type: DeviceType) -> int

def get_available_devices() -> List[Device]

import neurenix as nx
from neurenix.device import Device, DeviceType, get_device, get_available_devices

# Create a CPU device
cpu = Device(DeviceType.CPU)
print(cpu.name)  # "CPU"

# Create a CUDA device
cuda = Device(DeviceType.CUDA, index=0)
print(cuda.name)  # "CUDA:0"

# Create a TPU device
tpu = Device(DeviceType.TPU, index=0)

# Get device from string
device = get_device("cuda:1")
print(device.type)  # DeviceType.CUDA
print(device.index)  # 1

# Get available devices
devices = get_available_devices()
for dev in devices:
    print(f"Available: {dev.name}")

# Get device count
cuda_count = get_device_count(DeviceType.CUDA)
print(f"CUDA devices: {cuda_count}")

# Use device with tensors
tensor = nx.Tensor([1, 2, 3], device=cuda)
print(tensor.device.name)  # "CUDA:0"

# Move tensor to different device
tensor_cpu = tensor.to(cpu)
print(tensor_cpu.device.name)  # "CPU"

# Create WebGPU device for edge deployment
webgpu = Device(DeviceType.WEBGPU, index=0)

# Create NPU device for mobile AI
npu = Device(DeviceType.NPU, index=0)