update: test_sparse_supported

Author: kozistr

tests/test_gradients.py

@@ -43,39 +43,59 @@ def test_sparse_not_supported(no_sparse_optimizer):
         optimizer.step(lambda: 0.1)
 
 
+@pytest.mark.parametrize('sparse_optimizer', SPARSE_OPTIMIZERS)
+def test_sparse(sparse_optimizer):
+    opt = load_optimizer(optimizer=sparse_optimizer)
+
+    weight, weight_sparse = simple_sparse_parameter()
+
+    opt_dense = opt([weight], lr=1e-3, momentum=0.0)
+    opt_sparse = opt([weight_sparse], lr=1e-3, momentum=0.0)
+
+    opt_dense.step()
+    opt_sparse.step()
+    assert torch.allclose(weight, weight_sparse)
+
+    weight.grad = torch.rand_like(weight)
+    weight.grad[1] = 0.0
+    weight_sparse.grad = weight.grad.to_sparse()
+
+    opt_dense.step()
+    opt_sparse.step()
+    assert torch.allclose(weight, weight_sparse)
+
+    weight.grad = torch.rand_like(weight)
+    weight.grad[0] = 0.0
+    weight_sparse.grad = weight.grad.to_sparse()
+
+    opt_dense.step()
+    opt_sparse.step()
+    assert torch.allclose(weight, weight_sparse)
+
+
 @pytest.mark.parametrize('sparse_optimizer', SPARSE_OPTIMIZERS)
 def test_sparse_supported(sparse_optimizer):
     opt = load_optimizer(optimizer=sparse_optimizer)
 
-    optimizer = opt([simple_sparse_parameter()], momentum=0.0)
+    optimizer = opt([simple_sparse_parameter()[1]], momentum=0.0)
     optimizer.zero_grad()
     optimizer.step()
 
-    optimizer = opt([simple_sparse_parameter()], momentum=0.0)
-    with pytest.raises(RuntimeError):
-        optimizer.step()
-
-    optimizer = opt([simple_sparse_parameter()], momentum=0.0, eps=0.0)
-    optimizer.reset()
-    with pytest.raises(RuntimeError):
-        optimizer.step()
+    optimizer = opt([simple_sparse_parameter()[1]], momentum=0.0, eps=0.0)
+    optimizer.step()
 
     if sparse_optimizer == 'madgrad':
-        optimizer = opt([simple_sparse_parameter()], momentum=0.0, weight_decay=1e-3, decouple_decay=False)
-        optimizer.reset()
-
+        optimizer = opt([simple_sparse_parameter()[1]], momentum=0.0, weight_decay=1e-3, decouple_decay=False)
         with pytest.raises(NoSparseGradientError):
             optimizer.step()
 
-    optimizer = opt([simple_sparse_parameter()], momentum=0.9, weight_decay=1e-3)
+    optimizer = opt([simple_sparse_parameter()[1]], momentum=0.9, weight_decay=1e-3)
     optimizer.reset()
-
     if sparse_optimizer == 'madgrad':
         with pytest.raises(NoSparseGradientError):
             optimizer.step()
     else:
-        with pytest.raises(RuntimeError):
-            optimizer.step()
+        optimizer.step()
 
 
 @pytest.mark.parametrize('optimizer_name', VALID_OPTIMIZER_NAMES)