def test_gpu(self, coef, x, expected):
device = "cuda" if torch.cuda.is_available() else "cpu"
x = torch.as_tensor(x, dtype=torch.float, device=device)
x.requires_grad = True
coef = torch.as_tensor(coef, dtype=torch.float, device=device)
coef.requires_grad = True
result = polyval(coef, x)
if coef.shape[0] > 0: # empty coef doesn't have grad
result.mean().backward()
np.testing.assert_allclose(coef.grad.shape, coef.shape)
np.testing.assert_allclose(result.cpu().detach().numpy(), expected)
def test_floats(self, coef, x, expected):
result = polyval(coef, x)
np.testing.assert_allclose(result.cpu().numpy(), expected)