Repository: larq/larq

Issue:
Refactor layer test to use pytest parameterized test
Currently we use a mix of unittests written in `pytest` style and tests using `tf.test.TestCase` in larq. In #313 I added some simple `pytest` fixtures that allow to run tests in both eager and graph mode.

The only part of the test suite that uses `tf.test.TestCase` are the [layer tests](https://github.com/larq/larq/blob/f6f9277a006f6cb07b1e5f9bc591a087f2261152/larq/layers_test.py#L73-L187) which rely on [`keras_parameterized.run_all_keras_modes`](https://github.com/tensorflow/tensorflow/blob/669080135924d204a1c67dd556223d794efbf664/tensorflow/python/keras/keras_parameterized.py#L176-L303)  and a forked/copied version of [`testing_utils. layer_test`](https://github.com/tensorflow/tensorflow/blob/669080135924d204a1c67dd556223d794efbf664/tensorflow/python/keras/testing_utils.py#L75-L261).

I think it would be a lot cleaner to write a simple [`pytest` fixture](https://pytest.org/en/latest/fixture.html#fixture) similar to #313 that allows us to run a normal [`pytest` parameterized test](https://pytest.org/en/latest/parametrize.html#pytest-mark-parametrize) in all Keras modes and adapt the layer tests to use it. That way we can get rid of `absl.testing` and some requirements on TensorFlow internals.


Discussion:

Resolving patch:
diff --git a/larq/conftest.py b/larq/conftest.py
--- a/larq/conftest.py
+++ b/larq/conftest.py
@@ -1,4 +1,5 @@
 import pytest
+import tensorflow as tf
 from tensorflow.python.eager import context
 
 
@@ -21,3 +22,27 @@ def eager_and_graph_mode(request):
     """pytest fixture for running test in eager and graph mode"""
     with getattr(context, f"{request.param}_mode")():
         yield request.param
+
+
+@pytest.fixture(params=["graph", "tf_eager", "tf_keras_eager"])
+def keras_should_run_eagerly(request):
+    """Fixture to run in graph and two eager modes.
+
+    The modes are:
+    - Graph mode
+    - TensorFlow eager and Keras eager
+    - TensorFlow eager and Keras not eager
+
+    The `tf.context` sets graph/eager mode for TensorFlow. The yield is True if Keras
+    should run eagerly.
+    """
+
+    if request.param == "graph":
+        if int(tf.__version__[0]) >= 2:
+            pytest.skip("Skipping graph mode for TensorFlow 2+.")
+
+        with context.graph_mode():
+            yield
+    else:
+        with context.eager_mode():
+            yield request.param == "tf_keras_eager"
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -27,7 +27,6 @@ def readme():
         "tensorflow": ["tensorflow>=1.14.0"],
         "tensorflow_gpu": ["tensorflow-gpu>=1.14.0"],
         "test": [
-            "absl-py==0.8.1",
             "pytest==5.2.2",
             "pytest-cov==2.8.1",
             "pytest-xdist==1.30.0",


Test patch:
diff --git a/larq/layers_test.py b/larq/layers_test.py
--- a/larq/layers_test.py
+++ b/larq/layers_test.py
@@ -1,79 +1,78 @@
 import tensorflow as tf
 import numpy as np
-from absl.testing import parameterized
 import larq as lq
 import pytest
 import inspect
 from larq import testing_utils
-from tensorflow.python.keras import keras_parameterized
 
 
-def random_input(shape):
-    for i, dim in enumerate(shape):
-        if dim is None:
-            shape[i] = np.random.randint(1, 4)
-    data = 10 * np.random.random(shape) - 0.5
-    return data.astype("float32")
-
-
-parameterized_all_layers = parameterized.named_parameters(
-    ("QuantDense", lq.layers.QuantDense, tf.keras.layers.Dense, (3, 2), dict(units=3)),
+PARAMS_ALL_LAYERS = [
+    (lq.layers.QuantDense, tf.keras.layers.Dense, (3, 2), dict(units=3)),
     (
-        "QuantConv1D",
         lq.layers.QuantConv1D,
         tf.keras.layers.Conv1D,
         (2, 3, 7),
         dict(filters=2, kernel_size=3),
     ),
     (
-        "QuantConv2D",
         lq.layers.QuantConv2D,
         tf.keras.layers.Conv2D,
         (2, 3, 7, 6),
         dict(filters=2, kernel_size=3),
     ),
     (
-        "QuantConv3D",
         lq.layers.QuantConv3D,
         tf.keras.layers.Conv3D,
         (2, 3, 7, 6, 5),
         dict(filters=2, kernel_size=3),
     ),
     (
-        "QuantConv2DTranspose",
         lq.layers.QuantConv2DTranspose,
         tf.keras.layers.Conv2DTranspose,
         (2, 3, 7, 6),
         dict(filters=2, kernel_size=3),
     ),
     (
-        "QuantConv3DTranspose",
         lq.layers.QuantConv3DTranspose,
         tf.keras.layers.Conv3DTranspose,
         (2, 3, 7, 6, 5),
         dict(filters=2, kernel_size=3),
     ),
     (
-        "QuantLocallyConnected1D",
         lq.layers.QuantLocallyConnected1D,
         tf.keras.layers.LocallyConnected1D,
         (2, 8, 5),
         dict(filters=4, kernel_size=3),
     ),
     (
-        "QuantLocallyConnected2D",
         lq.layers.QuantLocallyConnected2D,
         tf.keras.layers.LocallyConnected2D,
         (8, 6, 10, 4),
         dict(filters=3, kernel_size=3),
     ),
-)
+]
+
+PARAMS_SEP_LAYERS = [
+    (lq.layers.QuantSeparableConv1D, tf.keras.layers.SeparableConv1D, (2, 3, 7),),
+    (lq.layers.QuantSeparableConv2D, tf.keras.layers.SeparableConv2D, (2, 3, 7, 6),),
+]
 
 
-@keras_parameterized.run_all_keras_modes
-class LayersTest(keras_parameterized.TestCase):
-    @parameterized_all_layers
-    def test_binarization(self, quantized_layer, layer, input_shape, kwargs):
+def random_input(shape):
+    for i, dim in enumerate(shape):
+        if dim is None:
+            shape[i] = np.random.randint(1, 4)
+    data = 10 * np.random.random(shape) - 0.5
+    return data.astype("float32")
+
+
+class TestLayers:
+    @pytest.mark.parametrize(
+        "quantized_layer, layer, input_shape, kwargs", PARAMS_ALL_LAYERS
+    )
+    def test_binarization(
+        self, quantized_layer, layer, input_shape, kwargs, keras_should_run_eagerly
+    ):
         input_data = random_input(input_shape)
         random_weight = np.random.random() - 0.5
 
@@ -87,6 +86,7 @@ def test_binarization(self, quantized_layer, layer, input_shape, kwargs):
                     kernel_initializer=tf.keras.initializers.constant(random_weight),
                 ),
                 input_data=input_data,
+                should_run_eagerly=keras_should_run_eagerly,
             )
 
         fp_output = testing_utils.layer_test(
@@ -98,54 +98,15 @@ def test_binarization(self, quantized_layer, layer, input_shape, kwargs):
                 ),
             ),
             input_data=np.sign(input_data),
+            should_run_eagerly=keras_should_run_eagerly,
         )
 
-        self.assertAllClose(quant_output, fp_output)
+        np.testing.assert_allclose(quant_output, fp_output)
 
-    def test_depthwise_layers(self):
-        input_data = random_input((2, 3, 7, 6))
-        random_weight = np.random.random() - 0.5
-
-        with lq.metrics.scope(["flip_ratio"]):
-            quant_output = testing_utils.layer_test(
-                lq.layers.QuantDepthwiseConv2D,
-                kwargs=dict(
-                    kernel_size=3,
-                    depthwise_quantizer="ste_sign",
-                    input_quantizer="ste_sign",
-                    depthwise_initializer=tf.keras.initializers.constant(random_weight),
-                ),
-                input_data=input_data,
-            )
-
-        fp_output = testing_utils.layer_test(
-            tf.keras.layers.DepthwiseConv2D,
-            kwargs=dict(
-                kernel_size=3,
-                depthwise_initializer=tf.keras.initializers.constant(
-                    np.sign(random_weight)
-                ),
-            ),
-            input_data=np.sign(input_data),
-        )
-
-        self.assertAllClose(quant_output, fp_output)
-
-    @parameterized.named_parameters(
-        (
-            "QuantSeparableConv1D",
-            lq.layers.QuantSeparableConv1D,
-            tf.keras.layers.SeparableConv1D,
-            (2, 3, 7),
-        ),
-        (
-            "QuantSeparableConv2D",
-            lq.layers.QuantSeparableConv2D,
-            tf.keras.layers.SeparableConv2D,
-            (2, 3, 7, 6),
-        ),
-    )
-    def test_separable_layers(self, quantized_layer, layer, input_shape):
+    @pytest.mark.parametrize("quantized_layer, layer, input_shape", PARAMS_SEP_LAYERS)
+    def test_separable_layers(
+        self, quantized_layer, layer, input_shape, keras_should_run_eagerly
+    ):
         input_data = random_input(input_shape)
         random_d_kernel = np.random.random() - 0.5
         random_p_kernel = np.random.random() - 0.5
@@ -167,6 +128,7 @@ def test_separable_layers(self, quantized_layer, layer, input_shape):
                     ),
                 ),
                 input_data=input_data,
+                should_run_eagerly=keras_should_run_eagerly,
             )
 
         fp_output = testing_utils.layer_test(
@@ -182,56 +144,84 @@ def test_separable_layers(self, quantized_layer, layer, input_shape):
                 ),
             ),
             input_data=np.sign(input_data),
+            should_run_eagerly=keras_should_run_eagerly,
         )
 
-        self.assertAllClose(quant_output, fp_output)
+        np.testing.assert_allclose(quant_output, fp_output)
 
+    def test_depthwise_layers(self, keras_should_run_eagerly):
+        input_data = random_input((2, 3, 7, 6))
+        random_weight = np.random.random() - 0.5
 
-def test_layer_warns(caplog):
-    lq.layers.QuantDense(5, kernel_quantizer="ste_sign")
-    assert len(caplog.records) >= 1
-    assert "kernel_constraint" in caplog.text
+        with lq.metrics.scope(["flip_ratio"]):
+            quant_output = testing_utils.layer_test(
+                lq.layers.QuantDepthwiseConv2D,
+                kwargs=dict(
+                    kernel_size=3,
+                    depthwise_quantizer="ste_sign",
+                    input_quantizer="ste_sign",
+                    depthwise_initializer=tf.keras.initializers.constant(random_weight),
+                ),
+                input_data=input_data,
+                should_run_eagerly=keras_should_run_eagerly,
+            )
 
+        fp_output = testing_utils.layer_test(
+            tf.keras.layers.DepthwiseConv2D,
+            kwargs=dict(
+                kernel_size=3,
+                depthwise_initializer=tf.keras.initializers.constant(
+                    np.sign(random_weight)
+                ),
+            ),
+            input_data=np.sign(input_data),
+            should_run_eagerly=keras_should_run_eagerly,
+        )
 
-def test_layer_does_not_warn(caplog):
-    lq.layers.QuantDense(
-        5, kernel_quantizer="ste_sign", kernel_constraint="weight_clip"
-    )
-    assert caplog.records == []
+        np.testing.assert_allclose(quant_output, fp_output)
 
 
-def test_depthwise_layer_warns(caplog):
-    lq.layers.QuantDepthwiseConv2D(5, depthwise_quantizer="ste_sign")
-    assert len(caplog.records) >= 1
-    assert "depthwise_constraint" in caplog.text
+class TestLayerWarns:
+    def test_layer_warns(self, caplog):
+        lq.layers.QuantDense(5, kernel_quantizer="ste_sign")
+        assert len(caplog.records) >= 1
+        assert "kernel_constraint" in caplog.text
 
+    def test_layer_does_not_warn(self, caplog):
+        lq.layers.QuantDense(
+            5, kernel_quantizer="ste_sign", kernel_constraint="weight_clip"
+        )
+        assert caplog.records == []
 
-def test_depthwise_layer_does_not_warn(caplog):
-    lq.layers.QuantDepthwiseConv2D(
-        5, depthwise_quantizer="ste_sign", depthwise_constraint="weight_clip"
-    )
-    assert caplog.records == []
+    def test_depthwise_layer_warns(self, caplog):
+        lq.layers.QuantDepthwiseConv2D(5, depthwise_quantizer="ste_sign")
+        assert len(caplog.records) >= 1
+        assert "depthwise_constraint" in caplog.text
 
+    def test_depthwise_layer_does_not_warn(self, caplog):
+        lq.layers.QuantDepthwiseConv2D(
+            5, depthwise_quantizer="ste_sign", depthwise_constraint="weight_clip"
+        )
+        assert caplog.records == []
 
-def test_separable_layer_warns(caplog):
-    lq.layers.QuantSeparableConv2D(
-        3, 3, depthwise_quantizer="ste_sign", pointwise_quantizer="ste_sign"
-    )
-    assert len(caplog.records) == 2
-    assert "depthwise_constraint" in caplog.text
-    assert "pointwise_constraint" in caplog.text
-
-
-def test_separable_layer_does_not_warn(caplog):
-    lq.layers.QuantSeparableConv2D(
-        3,
-        3,
-        depthwise_quantizer="ste_sign",
-        pointwise_quantizer="ste_sign",
-        depthwise_constraint="weight_clip",
-        pointwise_constraint="weight_clip",
-    )
-    assert caplog.records == []
+    def test_separable_layer_warns(self, caplog):
+        lq.layers.QuantSeparableConv2D(
+            3, 3, depthwise_quantizer="ste_sign", pointwise_quantizer="ste_sign"
+        )
+        assert len(caplog.records) == 2
+        assert "depthwise_constraint" in caplog.text
+        assert "pointwise_constraint" in caplog.text
+
+    def test_separable_layer_does_not_warn(self, caplog):
+        lq.layers.QuantSeparableConv2D(
+            3,
+            3,
+            depthwise_quantizer="ste_sign",
+            pointwise_quantizer="ste_sign",
+            depthwise_constraint="weight_clip",
+            pointwise_constraint="weight_clip",
+        )
+        assert caplog.records == []
 
 
 def test_metrics():
diff --git a/larq/testing_utils.py b/larq/testing_utils.py
--- a/larq/testing_utils.py
+++ b/larq/testing_utils.py
@@ -2,9 +2,6 @@
 import numpy as np
 import tensorflow as tf
 
-# We should find a better solution without relying on private objects
-from tensorflow.python.keras.testing_utils import _thread_local_data, should_run_eagerly
-
 
 def generate_real_values_with_zeros(low=-2, high=2, shape=(4, 10)):
     real_values = np.random.uniform(low, high, shape)
@@ -48,6 +45,7 @@ def layer_test(
     input_data=None,
     expected_output=None,
     expected_output_dtype=None,
+    should_run_eagerly=False,
 ):
     """Test routine for a layer with a single input and single output.
     Arguments:
@@ -148,17 +146,10 @@ def layer_test(
     # train(). This was causing some error for layer with Defun as it body.
     # See b/120160788 for more details. This should be mitigated after 2.0.
     model = tf.keras.models.Model(x, layer(x))
-    if _thread_local_data.run_eagerly is not None:
-        model.compile(
-            "rmsprop",
-            "mse",
-            weighted_metrics=["acc"],
-            run_eagerly=should_run_eagerly(),
-        )
-        model.train_on_batch(input_data, actual_output)
-    else:
-        model.compile("rmsprop", "mse", weighted_metrics=["acc"])
-        model.train_on_batch(input_data, actual_output)
+    model.compile(
+        "rmsprop", "mse", weighted_metrics=["acc"], run_eagerly=should_run_eagerly,
+    )
+    model.train_on_batch(input_data, actual_output)
 
     # test as first layer in Sequential API
     layer_config = layer.get_config()