Repository: tensorflow/addons

Issue:
CohenKappa raise exception for batch_size=1
**System information**
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Centos 7
- TensorFlow version and how it was installed (source or binary): 2.2.0
- TensorFlow-Addons version and how it was installed (source or binary): 0.10.0
- Python version: 3.6.9
- Is GPU used? (yes/no): no

**Describe the bug**

`tfa.metrics.CohenKappa` raise exception when batch_size=1.

**Code to reproduce the issue**

```python
import tensorflow as tf
import tensorflow_addons as tfa

kappa = tfa.metrics.CohenKappa(5, regression=True, sparse_labels=True)

# Batch-size > 1: this will not raise an exception
kappa.update_state(tf.ones(2), tf.ones(2))

# Batch-size = 1: this will raise an exception
kappa.update_state(tf.ones(1), tf.ones(1))
```

**Other info / logs**

The issue is due to `tf.squeeze` being use in `_update_confusion_matrix`:
https://github.com/tensorflow/addons/blob/5f618fdb92d9737da059de2a33fa606e97505398/tensorflow_addons/metrics/cohens_kappa.py#L170-L171

When the tensor has max shape equals to 1, like in `(1,)`, all dims are removed and a single scalar is returned by `tf.squeeze`. A subsequent call to `tf.math.confusion_matrix` will raise the exception while trying to stack tensors over an invalid axis.

```
ValueError                                Traceback (most recent call last)
<ipython-input-12-7ea8164c36b0> in <module>
      9
     10 # Batch-size = 1: this will raise an exception due to tf.squeeze
---> 11 kappa.update_state(tf.ones(1), tf.ones(1))

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/utils/metrics_utils.py in decorated(metric_obj, *args, **kwargs)
     88
     89     with tf_utils.graph_context_for_symbolic_tensors(*args, **kwargs):
---> 90       update_op = update_state_fn(*args, **kwargs)
     91     if update_op is not None:  # update_op will be None in eager execution.
     92       metric_obj.add_update(update_op)

/usr/local/lib/python3.6/dist-packages/tensorflow_addons/metrics/cohens_kappa.py in update_state(self, y_true, y_pred, sample_weight)
    137           Update op.
    138         """
--> 139         return self._update(y_true, y_pred, sample_weight)
    140
    141     def _update_binary_class_model(self, y_true, y_pred, sample_weight=None):

/usr/local/lib/python3.6/dist-packages/tensorflow_addons/metrics/cohens_kappa.py in _update_multi_class_model(self, y_true, y_pred, sample_weight)
    153         y_pred = self._cast_ypred(y_pred)
    154
--> 155         return self._update_confusion_matrix(y_true, y_pred, sample_weight)
    156
    157     @tf.function

/usr/local/lib/python3.6/dist-packages/tensorflow_addons/metrics/cohens_kappa.py in _update_confusion_matrix(self, y_true, y_pred, sample_weight)
    176             num_classes=self.num_classes,
    177             weights=sample_weight,
--> 178             dtype=tf.float32,
    179         )
    180

/usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/confusion_matrix.py in confusion_matrix(labels, predictions, num_classes, weights, dtype, name)
    190
    191     shape = array_ops.stack([num_classes, num_classes])
--> 192     indices = array_ops.stack([labels, predictions], axis=1)
    193     values = (array_ops.ones_like(predictions, dtype)
    194               if weights is None else weights)

/usr/local/lib/python3.6/dist-packages/tensorflow/python/util/dispatch.py in wrapper(*args, **kwargs)
    178     """Call target, and fall back on dispatchers if there is a TypeError."""
    179     try:
--> 180       return target(*args, **kwargs)
    181     except (TypeError, ValueError):
    182       # Note: convert_to_eager_tensor currently raises a ValueError, not a

/usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/array_ops.py in stack(values, axis, name)
   1338     if axis < -expanded_num_dims or axis >= expanded_num_dims:
   1339       raise ValueError("axis = %d not in [%d, %d)" %
-> 1340                        (axis, -expanded_num_dims, expanded_num_dims))
   1341
   1342   return gen_array_ops.pack(values, axis=axis, name=name)

ValueError: axis = 1 not in [-1, 1)
```



Discussion:

Resolving patch:
diff --git a/tensorflow_addons/metrics/cohens_kappa.py b/tensorflow_addons/metrics/cohens_kappa.py
--- a/tensorflow_addons/metrics/cohens_kappa.py
+++ b/tensorflow_addons/metrics/cohens_kappa.py
@@ -167,9 +167,19 @@ def _cast_ypred(self, y_pred):
             y_pred = tf.cast(y_pred, dtype=tf.int64)
         return y_pred
 
+    @tf.function
+    def _safe_squeeze(self, y):
+        y = tf.squeeze(y)
+
+        # Check for scalar result
+        if tf.rank(y) == 0:
+            y = tf.expand_dims(y, 0)
+
+        return y
+
     def _update_confusion_matrix(self, y_true, y_pred, sample_weight):
-        y_true = tf.squeeze(y_true)
-        y_pred = tf.squeeze(y_pred)
+        y_true = self._safe_squeeze(y_true)
+        y_pred = self._safe_squeeze(y_pred)
 
         new_conf_mtx = tf.math.confusion_matrix(
             labels=y_true,


Test patch:
diff --git a/tensorflow_addons/metrics/tests/cohens_kappa_test.py b/tensorflow_addons/metrics/tests/cohens_kappa_test.py
--- a/tensorflow_addons/metrics/tests/cohens_kappa_test.py
+++ b/tensorflow_addons/metrics/tests/cohens_kappa_test.py
@@ -13,7 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 """Tests for Cohen's Kappa Metric."""
-
 import pytest
 import numpy as np
 import tensorflow as tf
@@ -237,3 +236,13 @@ def test_cohen_kappa_serialization():
 
     ck = CohenKappa(num_classes=5, sparse_labels=True, weightage="quadratic")
     check_metric_serialization(ck, actuals, preds, weights)
+
+
+def test_cohen_kappa_single_batch():
+    # Test for issue #1962
+    obj = CohenKappa(5, regression=True, sparse_labels=True)
+
+    # Test single batch update
+    obj.update_state(tf.ones(1), tf.zeros(1))
+
+    np.testing.assert_allclose(0, obj.result().numpy())