Ensure indices list in sparse optimizer tests is unique

Summary:
There were no dimensionality constraints to the generated indices
array, causing many examples being generated and filtered out. Instead,
we should ensure the probability of unique indices is high.

There is a better fix for this by using the `unique` keyword argument
to `hypothesis.extra.numpy.arrays`, but this is available only in
hypothesis version 3.28.0 and later.

This is related to #1536 and #1599.

Once this change has proven to be OK, we can modify the other tests
that now have health check suppression enabled as well.
Closes https://github.com/caffe2/caffe2/pull/1686

Reviewed By: Yangqing

Differential Revision: D6651789

Pulled By: pietern

fbshipit-source-id: d80886c9ccf0a7a842a7580a279f33a2d6cca97c

3 files changed, 47 insertions, 18 deletions

diff --git a/caffe2/python/hypothesis_test_util.py b/caffe2/python/hypothesis_test_util.py
index 0260037d26..2f7da57ed5 100644
--- a/caffe2/python/hypothesis_test_util.py
+++ b/caffe2/python/hypothesis_test_util.py
@@ -133,12 +133,21 @@ def elements_of_type(dtype=np.float32, filter_=None):
 def arrays(dims, dtype=np.float32, elements=None):
     if elements is None:
         elements = elements_of_type(dtype)
-    return hypothesis.extra.numpy.arrays(dtype, dims, elements=elements)
+    return hypothesis.extra.numpy.arrays(
+        dtype,
+        dims,
+        elements=elements,
+    )
 
 
-def tensor(min_dim=1, max_dim=4, dtype=np.float32, elements=None, **kwargs):
+def tensor(min_dim=1,
+           max_dim=4,
+           dtype=np.float32,
+           elements=None,
+           **kwargs):
     dims_ = st.lists(dims(**kwargs), min_size=min_dim, max_size=max_dim)
-    return dims_.flatmap(lambda dims: arrays(dims, dtype, elements))
+    return dims_.flatmap(
+        lambda dims: arrays(dims, dtype, elements))
 
 
 def tensor1d(min_len=1, max_len=64, dtype=np.float32, elements=None):
@@ -243,8 +252,10 @@ def sparse_lengths_tensor(**kwargs):
 def tensors(n, min_dim=1, max_dim=4, dtype=np.float32, elements=None, **kwargs):
     dims_ = st.lists(dims(**kwargs), min_size=min_dim, max_size=max_dim)
     return dims_.flatmap(
-        lambda dims: st.lists(arrays(dims, dtype, elements),
-                              min_size=n, max_size=n))
+        lambda dims: st.lists(
+            arrays(dims, dtype, elements),
+            min_size=n,
+            max_size=n))
 
 
 def tensors1d(n, min_len=1, max_len=64, dtype=np.float32, elements=None):
diff --git a/caffe2/python/operator_test/adam_test.py b/caffe2/python/operator_test/adam_test.py
index 6e94f922be..0ac0a81ac0 100644..100755
--- a/caffe2/python/operator_test/adam_test.py
+++ b/caffe2/python/operator_test/adam_test.py
@@ -98,14 +98,20 @@ class TestAdam(hu.HypothesisTestCase):
 
         # Create an indexing array containing values which index into grad
         indices = data_strategy.draw(
-            hu.tensor(dtype=np.int64,
-                      elements=st.sampled_from(np.arange(grad.shape[0]))),
+            hu.tensor(
+                max_dim=1,
+                min_value=1,
+                max_value=grad.shape[0],
+                dtype=np.int64,
+                elements=st.sampled_from(np.arange(grad.shape[0])),
+            ),
         )
-        hypothesis.note('indices.shape: %s' % str(indices.shape))
 
-        # For now, the indices must be unique
-        hypothesis.assume(np.array_equal(np.unique(indices.flatten()),
-                                         np.sort(indices.flatten())))
+        # Verify that the generated indices are unique
+        hypothesis.assume(
+            np.array_equal(
+                np.unique(indices.flatten()),
+                np.sort(indices.flatten())))
 
         # Sparsify grad
         grad = grad[indices]
@@ -135,3 +141,8 @@ class TestAdam(hu.HypothesisTestCase):
             [param, mom1, mom2, indices, grad, LR, ITER],
             ref_sparse,
             input_device_options=input_device_options)
+
+
+if __name__ == "__main__":
+    import unittest
+    unittest.main()
diff --git a/caffe2/python/operator_test/momentum_sgd_test.py b/caffe2/python/operator_test/momentum_sgd_test.py
index 53b788e59a..9e178f5ff8 100644..100755
--- a/caffe2/python/operator_test/momentum_sgd_test.py
+++ b/caffe2/python/operator_test/momentum_sgd_test.py
@@ -100,23 +100,26 @@ class TestMomentumSGD(hu.HypothesisTestCase):
         # Create an indexing array containing values which index into grad
         indices = data_strategy.draw(
             hu.tensor(
+                max_dim=1,
+                min_value=1,
+                max_value=grad.shape[0],
                 dtype=np.int64,
-                elements=st.sampled_from(np.arange(grad.shape[0]))
+                elements=st.sampled_from(np.arange(grad.shape[0])),
             ),
         )
-        hypothesis.note('indices.shape: %s' % str(indices.shape))
 
-        # For now, the indices must be unique
+        # Verify that the generated indices are unique
         hypothesis.assume(
             np.array_equal(
-                np.unique(indices.flatten()), np.sort(indices.flatten())
-            )
-        )
+                np.unique(indices.flatten()),
+                np.sort(indices.flatten())))
 
         # Sparsify grad
         grad = grad[indices]
+
         # Make momentum >= 0
         m = np.abs(m)
+
         # Convert lr to a numpy array
         lr = np.asarray([lr], dtype=np.float32)
 
@@ -144,7 +147,11 @@ class TestMomentumSGD(hu.HypothesisTestCase):
             param[i] -= grad_new
             return (grad_new, m, param)
 
-        self.assertReferenceChecks(gc, op, [grad, m, lr, w, indices], sparse)
+        self.assertReferenceChecks(
+            gc,
+            op,
+            [grad, m, lr, w, indices],
+            sparse)
 
     @given(n=st.integers(4, 8), nesterov=st.booleans(), **hu.gcs_gpu_only)
     @unittest.skipIf(not workspace.has_gpu_support, "No gpu support.")