from graphgallery.functional import device
import tensorflow as tf
import torch


def test_device():
    # how about other backend?

    # tf
    assert isinstance(device("cpu", "tf"), str)
    assert device() == 'cpu'
    assert device("cpu", "tf") == 'CPU'
    assert device("cpu", "tf") == 'cpu'
    assert device("device/cpu", "tf") == 'cpu'

    try:
        assert device("gpu", "tf") == 'GPU'
        assert device("cuda", "tf") == 'GPU'
    except RuntimeError:
        pass
    device = tf.device("cpu")
    assert device(device, "tf") == device._device_name

    # ?? torch
    device = device("cpu", "torch")
    assert isinstance(device, torch.device) and 'cpu' in str(device)
    device = device(backend="torch")
    assert isinstance(device, torch.device) and 'cpu' in str(device)

    try:
        assert 'cuda' in str(device("gpu", "torch"))
        assert 'cuda' in str(device("cuda", "torch"))
    except RuntimeError:
        pass
    device = torch.device("cpu")
    assert device(device, "torch") == device
    
    
if __name__ == "__main__":
    test_device()    

[((496, 512), 'tensorflow.device', 'tf.device', (['"""cpu"""'], {}), "('cpu')\n", (505, 512), True, 'import tensorflow as tf\n'), ((597, 619), 'graphgallery.functional.device', 'device', (['"""cpu"""', '"""torch"""'], {}), "('cpu', 'torch')\n", (603, 619), False, 'from graphgallery.functional import device\n'), ((702, 725), 'graphgallery.functional.device', 'device', ([], {'backend': '"""torch"""'}), "(backend='torch')\n", (708, 725), False, 'from graphgallery.functional import device\n'), ((963, 982), 'torch.device', 'torch.device', (['"""cpu"""'], {}), "('cpu')\n", (975, 982), False, 'import torch\n'), ((164, 183), 'graphgallery.functional.device', 'device', (['"""cpu"""', '"""tf"""'], {}), "('cpu', 'tf')\n", (170, 183), False, 'from graphgallery.functional import device\n'), ((201, 209), 'graphgallery.functional.device', 'device', ([], {}), '()\n', (207, 209), False, 'from graphgallery.functional import device\n'), ((230, 249), 'graphgallery.functional.device', 'device', (['"""cpu"""', '"""tf"""'], {}), "('cpu', 'tf')\n", (236, 249), False, 'from graphgallery.functional import device\n'), ((270, 289), 'graphgallery.functional.device', 'device', (['"""cpu"""', '"""tf"""'], {}), "('cpu', 'tf')\n", (276, 289), False, 'from graphgallery.functional import device\n'), ((310, 336), 'graphgallery.functional.device', 'device', (['"""device/cpu"""', '"""tf"""'], {}), "('device/cpu', 'tf')\n", (316, 336), False, 'from graphgallery.functional import device\n'), ((524, 544), 'graphgallery.functional.device', 'device', (['device', '"""tf"""'], {}), "(device, 'tf')\n", (530, 544), False, 'from graphgallery.functional import device\n'), ((994, 1017), 'graphgallery.functional.device', 'device', (['device', '"""torch"""'], {}), "(device, 'torch')\n", (1000, 1017), False, 'from graphgallery.functional import device\n'), ((371, 390), 'graphgallery.functional.device', 'device', (['"""gpu"""', '"""tf"""'], {}), "('gpu', 'tf')\n", (377, 390), False, 'from graphgallery.functional import device\n'), ((415, 435), 'graphgallery.functional.device', 'device', (['"""cuda"""', '"""tf"""'], {}), "('cuda', 'tf')\n", (421, 435), False, 'from graphgallery.functional import device\n'), ((834, 856), 'graphgallery.functional.device', 'device', (['"""gpu"""', '"""torch"""'], {}), "('gpu', 'torch')\n", (840, 856), False, 'from graphgallery.functional import device\n'), ((887, 910), 'graphgallery.functional.device', 'device', (['"""cuda"""', '"""torch"""'], {}), "('cuda', 'torch')\n", (893, 910), False, 'from graphgallery.functional import device\n')]