# Buggy code
def _update_confusion_matrix(self, y_true, y_pred, sample_weight):
    y_true = tf.squeeze(y_true)
    y_pred = tf.squeeze(y_pred)

    new_conf_mtx = tf.math.confusion_matrix(
        labels=y_true,
        predictions=y_pred,
        num_classes=self.num_classes,
        weights=sample_weight,
        dtype=tf.float32,
    )

    return self.conf_mtx.assign_add(new_conf_mtx)


# Traceback (ValueError)
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)

# Fixed code
def _update_confusion_matrix(self, y_true, y_pred, sample_weight):
    y_true = self._safe_squeeze(y_true)
    y_pred = self._safe_squeeze(y_pred)

    new_conf_mtx = tf.math.confusion_matrix(
        labels=y_true,
        predictions=y_pred,
        num_classes=self.num_classes,
        weights=sample_weight,
        dtype=tf.float32,
    )

    return self.conf_mtx.assign_add(new_conf_mtx)