import unittest

import torch
import torch.nn as tnn
import torch.autograd as autograd

from common import gpu_test


class TestPyTorch(unittest.TestCase):
    # PyTorch smoke test based on http://pytorch.org/tutorials/beginner/nlp/deep_learning_tutorial.html
    def test_nn(self):
        torch.manual_seed(31337)
        linear_torch = tnn.Linear(5,3)
        data_torch = autograd.Variable(torch.randn(2, 5))
        linear_torch(data_torch)

    @gpu_test
    def test_linalg(self):
        A = torch.randn(3, 3).t().to('cuda')
        B = torch.randn(3).t().to('cuda')

        result = torch.linalg.solve(A, B)
        self.assertEqual(3, result.shape[0])

    @gpu_test
    def test_gpu_computation(self):
        cuda = torch.device('cuda')  
        a = torch.tensor([1., 2.], device=cuda)

        result = a.sum()

        self.assertEqual(torch.tensor([3.], device=cuda), result)

    @gpu_test
    def test_cuda_nn(self):
        # These throw if cuda is misconfigured
        tnn.GRUCell(10,10).cuda()
        tnn.RNNCell(10,10).cuda()
        tnn.LSTMCell(10,10).cuda()
        tnn.GRU(10,10).cuda()
        tnn.LSTM(10,10).cuda()
        tnn.RNN(10,10).cuda()