def get_xla_device_type(device: "torch.device") -> Optional[str]:
    """
    Returns the xla device type (CPU|GPU|TPU) or None if the device is a non-xla device.
    """
    if is_torch_tpu_available():
        return xm.xla_real_devices([device])[0].split(":")[0]
    return None

    Returns the xla device type (CPU|GPU|TPU) or None if the device is a non-xla device.
    
dpkd/transformers/src/transformers/training_args.py