-
Notifications
You must be signed in to change notification settings - Fork 1
Expand file tree
/
Copy pathtest_torchplot.py
More file actions
70 lines (58 loc) · 2.45 KB
/
test_torchplot.py
File metadata and controls
70 lines (58 loc) · 2.45 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
# Copyright The GeoML Team
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from collections import namedtuple
from inspect import getmembers, isfunction
import matplotlib.pyplot as plt
import numpy as np
import pytest
import torch
import torchplot as tp
Inputs = namedtuple("case", ["x", "y"])
_cases = [
Inputs(x=torch.randn(100), y=torch.randn(100)),
Inputs(x=torch.randn(100, requires_grad=True), y=torch.randn(100, requires_grad=True)),
# test that list/numpy arrays still works
Inputs(x=[1, 2, 3, 4], y=[1, 2, 3, 4]),
Inputs(x=np.random.randn(100), y=np.random.randn(100)),
# test that we can mix
Inputs(x=torch.randn(100), y=torch.randn(100, requires_grad=True)),
Inputs(x=np.random.randn(100), y=torch.randn(100, requires_grad=True)),
Inputs(x=torch.randn(5), y=[1, 2, 3, 4, 5]),
]
_members_to_check = [name for name, member in getmembers(plt) if isfunction(member) and not name.startswith("_")]
@pytest.mark.parametrize("member", _members_to_check)
def test_members(member):
""" test that all members have been copied """
assert member in dir(plt)
assert member in dir(tp)
@pytest.mark.parametrize("test_case", _cases)
def test_cpu(test_case):
""" test that it works on cpu """
# passed as args
assert tp.plot(test_case.x, test_case.y, ".")
# passed as kwargs
assert tp.scatter(x=test_case.x, y=test_case.y)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires cuda")
@pytest.mark.parametrize("test_case", _cases)
def test_gpu(test_case):
""" test that it works on gpu """
assert tp.plot(
test_case.x.cuda() if isinstance(test_case.x, torch.Tensor) else test_case.x,
test_case.y.cuda() if isinstance(test_case.y, torch.Tensor) else test_case.y,
)
# passed as kwargs
assert tp.scatter(
x=test_case.x.cuda() if isinstance(test_case.x, torch.Tensor) else test_case.x,
y=test_case.y.cuda() if isinstance(test_case.y, torch.Tensor) else test_case.y,
)