Repository: tensorflow/addons

Issue:
Unable to compile model using triplet loss 
**System information** 
- Have I written custom code (as opposed to using a stock example script provided in TensorFlow): custom
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Linux Ubuntu 16.04
- TensorFlow installed from (source or binary): binary
- TensorFlow version (use command below): tensoflow-gpu2beat1
- TensorFlow Addons installed from (source, PyPi): pip 0.4 
- TensorFlow Addons version: v0.4 
- Python version and type (eg. Anaconda Python, Stock Python as in Mac, or homebrew installed Python etc): Anaconda 
- Bazel version (if compiling from source): no
- GCC/Compiler version (if compiling from source): no
- Is GPU used? (yes/no): yes
- GPU model (if used):  GeForce GTX 1080 Ti

**Describe the bug**

 losses.triplet_semihard_loss from tensorflow addons does not compile successfully with keras model using tensorflow-gpu2beta1 and tensoflow addons v 0.4. 
It works with manual training but not compatible with model.compile and model.fit  

**Describe the expected behavior**
model should compile successfully. 
**Code to reproduce the issue**
```python
import tensorflow as tf
from tensorflow.keras.layers import  Dense
from tensorflow.keras.models import Model, Sequential
import tensorflow_addons as tfa
model = Sequential()
model.add(Dense(32, input_dim=784))
model.compile(loss=tfa.losses.triplet_semihard_loss, optimizer=tf.keras.optimizers.Adam())
```
Provide a reproducible test case that is the bare minimum necessary to generate the problem.

**Other info / logs**
---------------------------------------------------------------------------
```
AssertionError                            Traceback (most recent call last)
<ipython-input-13-f35ffd674f2d> in <module>
      5 model = Sequential()
      6 model.add(Dense(32, input_dim=784))
----> 7 model.compile(loss=tfa.losses.triplet_semihard_loss, optimizer=tf.keras.optimizers.Adam())

/anaconda3/envs/tf2/lib/python3.6/site-packages/tensorflow/python/training/tracking/base.py in _method_wrapper(self, *args, **kwargs)
    456     self._self_setattr_tracking = False  # pylint: disable=protected-access
    457     try:
--> 458       result = method(self, *args, **kwargs)
    459     finally:
    460       self._self_setattr_tracking = previous_value  # pylint: disable=protected-access

/anaconda3/envs/tf2/lib/python3.6/site-packages/tensorflow/python/keras/engine/training.py in compile(self, optimizer, loss, metrics, loss_weights, sample_weight_mode, weighted_metrics, target_tensors, distribute, **kwargs)
    335 
    336       # Creates the model loss and weighted metrics sub-graphs.
--> 337       self._compile_weights_loss_and_weighted_metrics()
    338 
    339       # Functions for train, test and predict will

/anaconda3/envs/tf2/lib/python3.6/site-packages/tensorflow/python/training/tracking/base.py in _method_wrapper(self, *args, **kwargs)
    456     self._self_setattr_tracking = False  # pylint: disable=protected-access
    457     try:
--> 458       result = method(self, *args, **kwargs)
    459     finally:
    460       self._self_setattr_tracking = previous_value  # pylint: disable=protected-access

/anaconda3/envs/tf2/lib/python3.6/site-packages/tensorflow/python/keras/engine/training.py in _compile_weights_loss_and_weighted_metrics(self, sample_weights)
   1492       #                   loss_weight_2 * output_2_loss_fn(...) +
   1493       #                   layer losses.
-> 1494       self.total_loss = self._prepare_total_loss(masks)
   1495 
   1496   def _prepare_skip_target_masks(self):

/anaconda3/envs/tf2/lib/python3.6/site-packages/tensorflow/python/keras/engine/training.py in _prepare_total_loss(self, masks)
   1552 
   1553           if hasattr(loss_fn, 'reduction'):
-> 1554             per_sample_losses = loss_fn.call(y_true, y_pred)
   1555             weighted_losses = losses_utils.compute_weighted_loss(
   1556                 per_sample_losses,

/anaconda3/envs/tf2/lib/python3.6/site-packages/tensorflow/python/keras/losses.py in call(self, y_true, y_pred)
    213       Loss values per sample.
    214     """
--> 215     return self.fn(y_true, y_pred, **self._fn_kwargs)
    216 
    217   def get_config(self):

/anaconda3/envs/tf2/lib/python3.6/site-packages/tensorflow/python/eager/def_function.py in __call__(self, *args, **kwds)
    406       # In this case we have not created variables on the first call. So we can
    407       # run the first trace but we should fail if variables are created.
--> 408       results = self._stateful_fn(*args, **kwds)
    409       if self._created_variables:
    410         raise ValueError("Creating variables on a non-first call to a function"

/anaconda3/envs/tf2/lib/python3.6/site-packages/tensorflow/python/eager/function.py in __call__(self, *args, **kwargs)
   1332   def __call__(self, *args, **kwargs):
   1333     """Calls a graph function specialized to the inputs."""
-> 1334     graph_function, args, kwargs = self._maybe_define_function(args, kwargs)
   1335     return graph_function._filtered_call(args, kwargs)  # pylint: disable=protected-access
   1336 

/anaconda3/envs/tf2/lib/python3.6/site-packages/tensorflow/python/eager/function.py in _maybe_define_function(self, args, kwargs)
   1646       graph_function = self._function_cache.primary.get(cache_key, None)
   1647       if graph_function is None:
-> 1648         graph_function = self._create_graph_function(args, kwargs)
   1649         self._function_cache.primary[cache_key] = graph_function
   1650       return graph_function, args, kwargs

/anaconda3/envs/tf2/lib/python3.6/site-packages/tensorflow/python/eager/function.py in _create_graph_function(self, args, kwargs, override_flat_arg_shapes)
   1539             arg_names=arg_names,
   1540             override_flat_arg_shapes=override_flat_arg_shapes,
-> 1541             capture_by_value=self._capture_by_value),
   1542         self._function_attributes)
   1543 

/anaconda3/envs/tf2/lib/python3.6/site-packages/tensorflow/python/framework/func_graph.py in func_graph_from_py_func(name, python_func, args, kwargs, signature, func_graph, autograph, autograph_options, add_control_dependencies, arg_names, op_return_value, collections, capture_by_value, override_flat_arg_shapes)
    714                                           converted_func)
    715 
--> 716       func_outputs = python_func(*func_args, **func_kwargs)
    717 
    718       # invariant: `func_outputs` contains only Tensors, CompositeTensors,

/anaconda3/envs/tf2/lib/python3.6/site-packages/tensorflow/python/eager/def_function.py in wrapped_fn(*args, **kwds)
    307         # __wrapped__ allows AutoGraph to swap in a converted function. We give
    308         # the function a weak reference to itself to avoid a reference cycle.
--> 309         return weak_wrapped_fn().__wrapped__(*args, **kwds)
    310     weak_wrapped_fn = weakref.ref(wrapped_fn)
    311 

/anaconda3/envs/tf2/lib/python3.6/site-packages/tensorflow/python/framework/func_graph.py in wrapper(*args, **kwargs)
    704           except Exception as e:  # pylint:disable=broad-except
    705             if hasattr(e, "ag_error_metadata"):
--> 706               raise e.ag_error_metadata.to_exception(type(e))
    707             else:
    708               raise

AssertionError: in converted code:

  /anaconda3/envs/tf2/lib/python3.6/site-packages/tensorflow_addons/losses/triplet.py:78 triplet_semihard_loss  *
        assert lshape.shape == 1

    AssertionError: 
```



Discussion:

Resolving patch:
diff --git a/tensorflow_addons/losses/triplet.py b/tensorflow_addons/losses/triplet.py
--- a/tensorflow_addons/losses/triplet.py
+++ b/tensorflow_addons/losses/triplet.py
@@ -73,9 +73,8 @@ def triplet_semihard_loss(y_true, y_pred, margin=1.0):
       margin: Float, margin term in the loss definition.
     """
     labels, embeddings = y_true, y_pred
-    # Reshape [batch_size] label tensor to a [batch_size, 1] label tensor.
+    # Reshape label tensor to [batch_size, 1].
     lshape = tf.shape(labels)
-    assert lshape.shape == 1
     labels = tf.reshape(labels, [lshape[0], 1])
 
     # Build pairwise squared distance matrix.


Test patch:
diff --git a/tensorflow_addons/losses/triplet_test.py b/tensorflow_addons/losses/triplet_test.py
--- a/tensorflow_addons/losses/triplet_test.py
+++ b/tensorflow_addons/losses/triplet_test.py
@@ -104,6 +104,13 @@ def test_unweighted(self):
         loss = cce_obj(y_true, y_pred)
         self.assertAlmostEqual(self.evaluate(loss), loss_np, 3)
 
+    def test_keras_model_compile(self):
+        model = tf.keras.models.Sequential([
+            tf.keras.layers.Input(shape=(784,)),
+            tf.keras.layers.Dense(10),
+        ])
+        model.compile(loss="triplet_semihard_loss", optimizer="adam")
+
 
 if __name__ == '__main__':
     tf.test.main()