from .context import aorun

import torch
from torch.autograd import Variable
from aorun.losses import mean_squared_error
from aorun.optimizers import SGD


def test_sgd_learning_rate():
    X = Variable(torch.rand(5, 3))
    y = Variable(torch.rand(5))
    w = Variable(torch.rand(3), requires_grad=True)
    opt = SGD([w], lr=0.1)

    o = X @ w
    loss = mean_squared_error(y, o)
    loss.backward()
    opt.step()

    assert loss > mean_squared_error(y, X @ w)


def test_sgd_momentum():
    X = Variable(torch.rand(5, 3))
    y = Variable(torch.rand(5))
    w = Variable(torch.rand(3), requires_grad=True)
    opt = SGD([w], lr=0.1, momentum=0.99)

    o = X @ w
    loss = mean_squared_error(y, o)
    loss.backward()
    opt.step()

    assert loss > mean_squared_error(y, X @ w)

@@ -1,12 +1,26 @@
+import pytest%0A
 from .contex
@@ -163,16 +163,130 @@
 t SGD%0A%0A%0A
+def test_sgd_without_params():%0A    opt = SGD(lr=0.1)%0A    with pytest.raises(Exception) as e:%0A        opt.step()%0A%0A%0A
 def test

improve optimizers test coverage