tensorflow: Using Precison metric in compile method raises shape mismatch error

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): No
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Windows 10 Home
• Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device: No mobile device
• TensorFlow installed from (source or binary): Installed using Pip command
• TensorFlow version (use command below): 2.1.0
• Python version: 3.6.8
• Bazel version (if compiling from source): No
• GCC/Compiler version (if compiling from source): No
• CUDA/cuDNN version: No
• GPU model and memory: No Describe the problem clearly here. Be sure to convey here why it’s a bug in TensorFlow or a feature request.

I am trying o implement different training metrics for keras sequential API. However when I try to implement precision method I get an error of shape mismatch. The same thing works when I use sigmoid as activation function instead of softmax.

I am trying to solve binary classification problem.

Code to create model:

def create_model():
    model = tf.keras.Sequential([
    feature_layer,
    tf.keras.layers.Dense(units = 12, activation='relu', use_bias = True, kernel_initializer= 'glorot_normal', bias_initializer = 'zeros', name = 'd1'),
    tf.keras.layers.Dense(units = 6, activation='relu', use_bias = True, kernel_initializer= 'glorot_normal', bias_initializer = 'zeros', name = 'd2'),
    tf.keras.layers.Dense(units = 2, activation='softmax', name = 'out')
    ])
    model.compile(optimizer='adam',
                  loss='sparse_categorical_crossentropy',
                  metrics=[tf.keras.metrics.Precision()])
    return model

The same code runs when I try to run with sigmoid activation fuction with 1 output unit and Binary Crossentropy as my loss.

Source code / logs

Traceback (most recent call last):
  File "4.py", line 217, in <module>
    callbacks = ALL_CALLBACKS)
  File "C:\Users\Aniket\Desktop\Class\class_env\lib\site-packages\tensorflow_core\python\keras\engine\training.py", line 819, in fit
    use_multiprocessing=use_multiprocessing)
  File "C:\Users\Aniket\Desktop\Class\class_env\lib\site-packages\tensorflow_core\python\keras\engine\training_v2.py", line 342, in fit
    total_epochs=epochs)
  File "C:\Users\Aniket\Desktop\Class\class_env\lib\site-packages\tensorflow_core\python\keras\engine\training_v2.py", line 128, in run_one_epoch
    batch_outs = execution_function(iterator)
  File "C:\Users\Aniket\Desktop\Class\class_env\lib\site-packages\tensorflow_core\python\keras\engine\training_v2_utils.py", line 98, in execution_function
    distributed_function(input_fn))
  File "C:\Users\Aniket\Desktop\Class\class_env\lib\site-packages\tensorflow_core\python\eager\def_function.py", line 568, in __call__
    result = self._call(*args, **kwds)
  File "C:\Users\Aniket\Desktop\Class\class_env\lib\site-packages\tensorflow_core\python\eager\def_function.py", line 615, in _call
    self._initialize(args, kwds, add_initializers_to=initializers)
  File "C:\Users\Aniket\Desktop\Class\class_env\lib\site-packages\tensorflow_core\python\eager\def_function.py", line 497, in _initialize
    *args, **kwds))
  File "C:\Users\Aniket\Desktop\Class\class_env\lib\site-packages\tensorflow_core\python\eager\function.py", line 2389, in _get_concrete_function_internal_garbage_collected
    graph_function, _, _ = self._maybe_define_function(args, kwargs)
  File "C:\Users\Aniket\Desktop\Class\class_env\lib\site-packages\tensorflow_core\python\eager\function.py", line 2703, in _maybe_define_function
    graph_function = self._create_graph_function(args, kwargs)
  File "C:\Users\Aniket\Desktop\Class\class_env\lib\site-packages\tensorflow_core\python\eager\function.py", line 2593, in _create_graph_function
    capture_by_value=self._capture_by_value),
  File "C:\Users\Aniket\Desktop\Class\class_env\lib\site-packages\tensorflow_core\python\framework\func_graph.py", line 978, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "C:\Users\Aniket\Desktop\Class\class_env\lib\site-packages\tensorflow_core\python\eager\def_function.py", line 439, in wrapped_fn
    return weak_wrapped_fn().__wrapped__(*args, **kwds)
  File "C:\Users\Aniket\Desktop\Class\class_env\lib\site-packages\tensorflow_core\python\keras\engine\training_v2_utils.py", line 85, in distributed_function
    per_replica_function, args=args)
  File "C:\Users\Aniket\Desktop\Class\class_env\lib\site-packages\tensorflow_core\python\distribute\distribute_lib.py", line 763, in experimental_run_v2
    return self._extended.call_for_each_replica(fn, args=args, kwargs=kwargs)
  File "C:\Users\Aniket\Desktop\Class\class_env\lib\site-packages\tensorflow_core\python\distribute\distribute_lib.py", line 1819, in call_for_each_replica
    return self._call_for_each_replica(fn, args, kwargs)
  File "C:\Users\Aniket\Desktop\Class\class_env\lib\site-packages\tensorflow_core\python\distribute\distribute_lib.py", line 2164, in _call_for_each_replica
    return fn(*args, **kwargs)
  File "C:\Users\Aniket\Desktop\Class\class_env\lib\site-packages\tensorflow_core\python\autograph\impl\api.py", line 292, in wrapper
    return func(*args, **kwargs)
  File "C:\Users\Aniket\Desktop\Class\class_env\lib\site-packages\tensorflow_core\python\keras\engine\training_v2_utils.py", line 433, in train_on_batch
    output_loss_metrics=model._output_loss_metrics)
  File "C:\Users\Aniket\Desktop\Class\class_env\lib\site-packages\tensorflow_core\python\keras\engine\training_eager.py", line 316, in train_on_batch
    model, outs, targets, sample_weights=sample_weights, masks=masks)
  File "C:\Users\Aniket\Desktop\Class\class_env\lib\site-packages\tensorflow_core\python\keras\engine\training_eager.py", line 74, in _eager_metrics_fn
    skip_target_masks=model._prepare_skip_target_masks())
  File "C:\Users\Aniket\Desktop\Class\class_env\lib\site-packages\tensorflow_core\python\keras\engine\training.py", line 2004, in _handle_metrics
    target, output, output_mask))
  File "C:\Users\Aniket\Desktop\Class\class_env\lib\site-packages\tensorflow_core\python\keras\engine\training.py", line 1955, in _handle_per_output_metrics
    metric_fn, y_true, y_pred, weights=weights, mask=mask)
  File "C:\Users\Aniket\Desktop\Class\class_env\lib\site-packages\tensorflow_core\python\keras\engine\training_utils.py", line 1155, in call_metric_function
    return metric_fn(y_true, y_pred, sample_weight=weights)
  File "C:\Users\Aniket\Desktop\Class\class_env\lib\site-packages\tensorflow_core\python\keras\metrics.py", line 196, in __call__
    replica_local_fn, *args, **kwargs)
  File "C:\Users\Aniket\Desktop\Class\class_env\lib\site-packages\tensorflow_core\python\keras\distribute\distributed_training_utils.py", line 1135, in call_replica_local_fn
    return fn(*args, **kwargs)
  File "C:\Users\Aniket\Desktop\Class\class_env\lib\site-packages\tensorflow_core\python\keras\metrics.py", line 179, in replica_local_fn
    update_op = self.update_state(*args, **kwargs)  # pylint: disable=not-callable
  File "C:\Users\Aniket\Desktop\Class\class_env\lib\site-packages\tensorflow_core\python\keras\utils\metrics_utils.py", line 76, in decorated
    update_op = update_state_fn(*args, **kwargs)
  File "C:\Users\Aniket\Desktop\Class\class_env\lib\site-packages\tensorflow_core\python\keras\metrics.py", line 1216, in update_state
    sample_weight=sample_weight)
  File "C:\Users\Aniket\Desktop\Class\class_env\lib\site-packages\tensorflow_core\python\keras\utils\metrics_utils.py", line 303, in update_confusion_matrix_variables
    y_pred.shape.assert_is_compatible_with(y_true.shape)
  File "C:\Users\Aniket\Desktop\Class\class_env\lib\site-packages\tensorflow_core\python\framework\tensor_shape.py", line 1110, in assert_is_compatible_with
    raise ValueError("Shapes %s and %s are incompatible" % (self, other))
ValueError: Shapes (None, 2) and (None, 1) are incompatible