Repository: tensorflow/addons

Issue:
quantile_test.py is not run at all with bazel 
See #1201

**Describe the bug**

One of the test doesn't work, but we couldn't see it because bazel didn't run it.

**Code to reproduce the issue**

See #1201

**Other info / logs**

Logs when actually running the tests in the file:

```
(proof_of_concept_pytest) /mnt/c/Users/gdemarmi/Desktop/projects/addons $ pytest -v tensorflow_addons/losses/quantiles_test.py
/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/pep8.py:110: FutureWarning: Possible nested set at position 1
  EXTRANEOUS_WHITESPACE_REGEX = re.compile(r'[[({] | []}),;:]')
================================================================= test session starts ==================================================================
platform linux -- Python 3.7.5, pytest-5.0.1, py-1.8.0, pluggy-0.13.0 -- /home/gdemarmi/softwares/python/anaconda/bin/python
cachedir: .pytest_cache
rootdir: /mnt/c/Users/gdemarmi/Desktop/projects/addons
plugins: arraydiff-0.3, cov-2.7.1, doctestplus-0.2.0, flake8-1.0.4, forked-1.0.2, openfiles-0.4.0, pep8-1.0.6, remotedata-0.3.2, xdist-1.27.0, typeguard-2.7.1
collected 15 items

tensorflow_addons/losses/quantiles_test.py::PinballLossTest::test_all_correct_unweighted PASSED                                                  [  6%]
tensorflow_addons/losses/quantiles_test.py::PinballLossTest::test_config PASSED                                                                  [ 13%]
tensorflow_addons/losses/quantiles_test.py::PinballLossTest::test_invalid_sample_weight FAILED                                                   [ 20%]
tensorflow_addons/losses/quantiles_test.py::PinballLossTest::test_no_reduction PASSED                                                            [ 26%]
tensorflow_addons/losses/quantiles_test.py::PinballLossTest::test_sample_weighted PASSED                                                         [ 33%]
tensorflow_addons/losses/quantiles_test.py::PinballLossTest::test_scalar_weighted PASSED                                                         [ 40%]
tensorflow_addons/losses/quantiles_test.py::PinballLossTest::test_session SKIPPED                                                                [ 46%]
tensorflow_addons/losses/quantiles_test.py::PinballLossTest::test_sum_reduction PASSED                                                           [ 53%]
tensorflow_addons/losses/quantiles_test.py::PinballLossTest::test_timestep_weighted FAILED                                                       [ 60%]
tensorflow_addons/losses/quantiles_test.py::PinballLossTest::test_unweighted PASSED                                                              [ 66%]
tensorflow_addons/losses/quantiles_test.py::PinballLossTest::test_unweighted_quantile_0pc PASSED                                                 [ 73%]
tensorflow_addons/losses/quantiles_test.py::PinballLossTest::test_unweighted_quantile_100pc PASSED                                               [ 80%]
tensorflow_addons/losses/quantiles_test.py::PinballLossTest::test_unweighted_quantile_10pc PASSED                                                [ 86%]
tensorflow_addons/losses/quantiles_test.py::PinballLossTest::test_unweighted_quantile_90pc PASSED                                                [ 93%]
tensorflow_addons/losses/quantiles_test.py::PinballLossTest::test_zero_weighted PASSED                                                           [100%]

======================================================================= FAILURES =======================================================================
______________________________________________________ PinballLossTest.test_invalid_sample_weight ______________________________________________________

    def _create_c_op(graph, node_def, inputs, control_inputs):
      """Creates a TF_Operation.

      Args:
        graph: a `Graph`.
        node_def: `node_def_pb2.NodeDef` for the operation to create.
        inputs: A list of `Tensor`s (corresponding to scalar inputs) and lists of
          `Tensor`s (corresponding to sequence inputs, e.g. "int64 * N",
          "list(int64)"). The length of the list should be equal to the number of
          inputs specified by this operation's op def.
        control_inputs: A list of `Operation`s to set as control dependencies.

      Returns:
        A wrapped TF_Operation*.
      """
      # pylint: disable=protected-access
      op_desc = c_api.TF_NewOperation(graph._c_graph, compat.as_str(node_def.op),
                                      compat.as_str(node_def.name))
      if node_def.device:
        c_api.TF_SetDevice(op_desc, compat.as_str(node_def.device))
      # Add inputs
      for op_input in inputs:
        if isinstance(op_input, (list, tuple)):
          c_api.TF_AddInputList(op_desc, [t._as_tf_output() for t in op_input])
        else:
          c_api.TF_AddInput(op_desc, op_input._as_tf_output())

      # Add control inputs
      for control_input in control_inputs:
        c_api.TF_AddControlInput(op_desc, control_input._c_op)
      # pylint: enable=protected-access

      # Add attrs
      for name, attr_value in node_def.attr.items():
        serialized = attr_value.SerializeToString()
        # TODO(skyewm): this creates and deletes a new TF_Status for every attr.
        # It might be worth creating a convenient way to re-use the same status.
        c_api.TF_SetAttrValueProto(op_desc, compat.as_str(name), serialized)

      try:
>       c_op = c_api.TF_FinishOperation(op_desc)
E       tensorflow.python.framework.errors_impl.InvalidArgumentError: Can not squeeze dim[1], expected a dimension of 1, got 2 for 'pinball_loss/weighted_loss/Squeeze' (op: 'Squeeze') with input shapes: [2,2].

/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py:1619: InvalidArgumentError

During handling of the above exception, another exception occurred:
Traceback (most recent call last):
  File "/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py", line 1619, in _create_c_op
    c_op = c_api.TF_FinishOperation(op_desc)
tensorflow.python.framework.errors_impl.InvalidArgumentError: Can not squeeze dim[1], expected a dimension of 1, got 2 for 'pinball_loss/weighted_loss/Squeeze' (op: 'Squeeze') with input shapes: [2,2].

During handling of the above exception, another exception occurred:

ValueError: Can not squeeze dim[1], expected a dimension of 1, got 2 for 'pinball_loss/weighted_loss/Squeeze' (op: 'Squeeze') with input shapes: [2,2].

During handling of the above exception, another exception occurred:

self = <tensorflow_addons.losses.quantiles_test.PinballLossTest testMethod=test_invalid_sample_weight>, args = (), kwargs = {}

    def decorated(self, *args, **kwargs):
      try:
        with context.graph_mode():
          with self.test_session(use_gpu=use_gpu, config=config):
>           f(self, *args, **kwargs)

/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/tensorflow_core/python/framework/test_util.py:1111:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
tensorflow_addons/losses/quantiles_test.py:113: in test_invalid_sample_weight
    pin_obj(y_true, y_pred, sample_weight=sample_weight)
E   AssertionError: "weights can not be broadcast to values" does not match "Can not squeeze dim[1], expected a dimension of 1, got 2 for 'pinball_loss/weighted_loss/Squeeze' (op: 'Squeeze') with input shapes: [2,2]."
________________________________________________________ PinballLossTest.test_timestep_weighted ________________________________________________________

graph = <tensorflow.python.framework.ops.Graph object at 0x7f562c26e950>
node_def = name: "pinball_loss/weighted_loss/Squeeze"
op: "Squeeze"
attr {
  key: "T"
  value {
    type: DT_FLOAT
  }
}
attr {
  key: "squeeze_dims"
  value {
    list {
      i: -1
    }
  }
}

inputs = [<tf.Tensor 'pinball_loss/weighted_loss/Cast:0' shape=(2, 3) dtype=float32>], control_inputs = []

    def _create_c_op(graph, node_def, inputs, control_inputs):
      """Creates a TF_Operation.

      Args:
        graph: a `Graph`.
        node_def: `node_def_pb2.NodeDef` for the operation to create.
        inputs: A list of `Tensor`s (corresponding to scalar inputs) and lists of
          `Tensor`s (corresponding to sequence inputs, e.g. "int64 * N",
          "list(int64)"). The length of the list should be equal to the number of
          inputs specified by this operation's op def.
        control_inputs: A list of `Operation`s to set as control dependencies.

      Returns:
        A wrapped TF_Operation*.
      """
      # pylint: disable=protected-access
      op_desc = c_api.TF_NewOperation(graph._c_graph, compat.as_str(node_def.op),
                                      compat.as_str(node_def.name))
      if node_def.device:
        c_api.TF_SetDevice(op_desc, compat.as_str(node_def.device))
      # Add inputs
      for op_input in inputs:
        if isinstance(op_input, (list, tuple)):
          c_api.TF_AddInputList(op_desc, [t._as_tf_output() for t in op_input])
        else:
          c_api.TF_AddInput(op_desc, op_input._as_tf_output())

      # Add control inputs
      for control_input in control_inputs:
        c_api.TF_AddControlInput(op_desc, control_input._c_op)
      # pylint: enable=protected-access

      # Add attrs
      for name, attr_value in node_def.attr.items():
        serialized = attr_value.SerializeToString()
        # TODO(skyewm): this creates and deletes a new TF_Status for every attr.
        # It might be worth creating a convenient way to re-use the same status.
        c_api.TF_SetAttrValueProto(op_desc, compat.as_str(name), serialized)

      try:
>       c_op = c_api.TF_FinishOperation(op_desc)
E       tensorflow.python.framework.errors_impl.InvalidArgumentError: Can not squeeze dim[1], expected a dimension of 1, got 3 for 'pinball_loss/weighted_loss/Squeeze' (op: 'Squeeze') with input shapes: [2,3].

/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py:1619: InvalidArgumentError

During handling of the above exception, another exception occurred:

self = <tensorflow_addons.losses.quantiles_test.PinballLossTest testMethod=test_timestep_weighted>, args = (), kwargs = {}

    def decorated(self, *args, **kwargs):
      try:
        with context.graph_mode():
          with self.test_session(use_gpu=use_gpu, config=config):
>           f(self, *args, **kwargs)

/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/tensorflow_core/python/framework/test_util.py:1111:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
tensorflow_addons/losses/quantiles_test.py:95: in test_timestep_weighted
    loss = pin_obj(y_true, y_pred, sample_weight=sample_weight)
/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/tensorflow_core/python/keras/losses.py:128: in __call__
    losses, sample_weight, reduction=self._get_reduction())
/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/tensorflow_core/python/keras/utils/losses_utils.py:107: in compute_weighted_loss
    losses, sample_weight)
/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/tensorflow_core/python/ops/losses/util.py:145: in scale_losses_by_sample_weight
    losses, None, sample_weight)
/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/tensorflow_core/python/ops/losses/util.py:97: in squeeze_or_expand_dimensions
    sample_weight = array_ops.squeeze(sample_weight, [-1])
/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/tensorflow_core/python/util/dispatch.py:180: in wrapper
    return target(*args, **kwargs)
/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/tensorflow_core/python/util/deprecation.py:507: in new_func
    return func(*args, **kwargs)
/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/tensorflow_core/python/ops/array_ops.py:3780: in squeeze
    return gen_array_ops.squeeze(input, axis, name)
/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/tensorflow_core/python/ops/gen_array_ops.py:9231: in squeeze
    "Squeeze", input=input, squeeze_dims=axis, name=name)
/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/tensorflow_core/python/framework/op_def_library.py:742: in _apply_op_helper
    attrs=attr_protos, op_def=op_def)
/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py:3322: in _create_op_internal
    op_def=op_def)
/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py:1786: in __init__
    control_input_ops)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

graph = <tensorflow.python.framework.ops.Graph object at 0x7f562c26e950>
node_def = name: "pinball_loss/weighted_loss/Squeeze"
op: "Squeeze"
attr {
  key: "T"
  value {
    type: DT_FLOAT
  }
}
attr {
  key: "squeeze_dims"
  value {
    list {
      i: -1
    }
  }
}

inputs = [<tf.Tensor 'pinball_loss/weighted_loss/Cast:0' shape=(2, 3) dtype=float32>], control_inputs = []

    def _create_c_op(graph, node_def, inputs, control_inputs):
      """Creates a TF_Operation.

      Args:
        graph: a `Graph`.
        node_def: `node_def_pb2.NodeDef` for the operation to create.
        inputs: A list of `Tensor`s (corresponding to scalar inputs) and lists of
          `Tensor`s (corresponding to sequence inputs, e.g. "int64 * N",
          "list(int64)"). The length of the list should be equal to the number of
          inputs specified by this operation's op def.
        control_inputs: A list of `Operation`s to set as control dependencies.

      Returns:
        A wrapped TF_Operation*.
      """
      # pylint: disable=protected-access
      op_desc = c_api.TF_NewOperation(graph._c_graph, compat.as_str(node_def.op),
                                      compat.as_str(node_def.name))
      if node_def.device:
        c_api.TF_SetDevice(op_desc, compat.as_str(node_def.device))
      # Add inputs
      for op_input in inputs:
        if isinstance(op_input, (list, tuple)):
          c_api.TF_AddInputList(op_desc, [t._as_tf_output() for t in op_input])
        else:
          c_api.TF_AddInput(op_desc, op_input._as_tf_output())

      # Add control inputs
      for control_input in control_inputs:
        c_api.TF_AddControlInput(op_desc, control_input._c_op)
      # pylint: enable=protected-access

      # Add attrs
      for name, attr_value in node_def.attr.items():
        serialized = attr_value.SerializeToString()
        # TODO(skyewm): this creates and deletes a new TF_Status for every attr.
        # It might be worth creating a convenient way to re-use the same status.
        c_api.TF_SetAttrValueProto(op_desc, compat.as_str(name), serialized)

      try:
        c_op = c_api.TF_FinishOperation(op_desc)
      except errors.InvalidArgumentError as e:
        # Convert to ValueError for backwards compatibility.
>       raise ValueError(str(e))
E       ValueError: Can not squeeze dim[1], expected a dimension of 1, got 3 for 'pinball_loss/weighted_loss/Squeeze' (op: 'Squeeze') with input shapes: [2,3].

/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py:1622: ValueError
----------------------------------------------------------------- Captured stderr call -----------------------------------------------------------------
W0302 15:37:30.724358 140009831860032 def_function.py:586] 8 out of the last 11 calls to <function pinball_loss at 0x7f562f7ecb00> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings is likely due to passing python objects instead of tensors. Also, tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. Please refer to https://www.tensorflow.org/tutorials/customization/performance#python_or_tensor_args and https://www.tensorflow.org/api_docs/python/tf/function for more details.
------------------------------------------------------------------ Captured log call -------------------------------------------------------------------
WARNING  tensorflow:def_function.py:586 8 out of the last 11 calls to <function pinball_loss at 0x7f562f7ecb00> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings is likely due to passing python objects instead of tensors. Also, tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. Please refer to https://www.tensorflow.org/tutorials/customization/performance#python_or_tensor_args and https://www.tensorflow.org/api_docs/python/tf/function for more details.
=================================================================== warnings summary ===================================================================
/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/google/protobuf/descriptor.py:47
  /home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/google/protobuf/descriptor.py:47: DeprecationWarning: Using or importing the ABCs from 'collections' instead of from 'collections.abc' is deprecated since Python 3.3,and in 3.9 it will stop working
    from google.protobuf.pyext import _message

/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/botocore/vendored/requests/packages/urllib3/_collections.py:1
/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/botocore/vendored/requests/packages/urllib3/_collections.py:1
  /home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/botocore/vendored/requests/packages/urllib3/_collections.py:1: DeprecationWarning: Using or importing the ABCs from 'collections' instead of from 'collections.abc' is deprecated since Python 3.3,and in 3.9 it will stop working
    from collections import Mapping, MutableMapping

-- Docs: https://docs.pytest.org/en/latest/warnings.html
============================================= 2 failed, 12 passed, 1 skipped, 3 warnings in 17.15 seconds ==============================================
```



Discussion:

Resolving patch:
diff --git a/tensorflow_addons/losses/quantiles.py b/tensorflow_addons/losses/quantiles.py
--- a/tensorflow_addons/losses/quantiles.py
+++ b/tensorflow_addons/losses/quantiles.py
@@ -19,8 +19,8 @@
 from tensorflow_addons.utils.types import TensorLike, FloatTensorLike
 
 
-@tf.keras.utils.register_keras_serializable(package="Addons")
 @tf.function
+@tf.keras.utils.register_keras_serializable(package="Addons")
 def pinball_loss(
     y_true: TensorLike, y_pred: TensorLike, tau: FloatTensorLike = 0.5
 ) -> tf.Tensor:
@@ -62,14 +62,13 @@ def pinball_loss(
     y_pred = tf.convert_to_tensor(y_pred)
     y_true = tf.cast(y_true, y_pred.dtype)
 
-    # broadcast the pinball slope along the batch dimension, and clip to
-    # acceptable values
+    # Broadcast the pinball slope along the batch dimension
     tau = tf.expand_dims(tf.cast(tau, y_pred.dtype), 0)
     one = tf.cast(1, tau.dtype)
 
     delta_y = y_true - y_pred
     pinball = tf.math.maximum(tau * delta_y, (tau - one) * delta_y)
-    return tf.reduce_mean(tf.keras.backend.batch_flatten(pinball), axis=-1)
+    return tf.reduce_mean(pinball, axis=-1)
 
 
 @tf.keras.utils.register_keras_serializable(package="Addons")


Test patch:
diff --git a/tensorflow_addons/losses/tests/quantiles_test.py b/tensorflow_addons/losses/tests/quantiles_test.py
--- a/tensorflow_addons/losses/tests/quantiles_test.py
+++ b/tensorflow_addons/losses/tests/quantiles_test.py
@@ -20,6 +20,7 @@
 import tensorflow as tf
 
 from tensorflow_addons.losses import quantiles
+from distutils.version import LooseVersion
 
 
 def test_config():
@@ -63,7 +64,6 @@ def test_sample_weighted():
 
 
 def test_timestep_weighted():
-    pytest.skip("Failing. See https://github.com/tensorflow/addons/issues/1202")
     pin_obj = quantiles.PinballLoss()
     y_true = tf.constant([1, 9, 2, -5, -2, 6], shape=(2, 3, 1))
     y_pred = tf.constant([4, 8, 12, 8, 1, 3], shape=(2, 3, 1), dtype=tf.dtypes.float32)
@@ -81,13 +81,16 @@ def test_zero_weighted():
 
 
 def test_invalid_sample_weight():
-    pytest.skip("Failing. See https://github.com/tensorflow/addons/issues/1202")
     pin_obj = quantiles.PinballLoss()
     y_true = tf.constant([1, 9, 2, -5, -2, 6], shape=(2, 3, 1))
     y_pred = tf.constant([4, 8, 12, 8, 1, 3], shape=(2, 3, 1))
     sample_weight = tf.constant([3, 6, 5, 0], shape=(2, 2))
-    with pytest.raises(ValueError, match="weights can not be broadcast to values"):
-        pin_obj(y_true, y_pred, sample_weight=sample_weight)
+    if LooseVersion(tf.__version__) >= "2.2":
+        with pytest.raises(tf.errors.InvalidArgumentError, match="Incompatible shapes"):
+            pin_obj(y_true, y_pred, sample_weight=sample_weight)
+    else:
+        with pytest.raises(ValueError, match="weights can not be broadcast to values"):
+            pin_obj(y_true, y_pred, sample_weight=sample_weight)
 
 
 def test_unweighted_quantile_0pc():