# Buggy code
def triplet_semihard_loss(y_true, y_pred, margin=1.0):
    """Computes the triplet loss with semi-hard negative mining.

    Args:
      y_true: 1-D integer `Tensor` with shape [batch_size] of
        multiclass integer labels.
      y_pred: 2-D float `Tensor` of embedding vectors. Embeddings should
        be l2 normalized.
      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.
    lshape = tf.shape(labels)
    assert lshape.shape == 1
    labels = tf.reshape(labels, [lshape[0], 1])

    # Build pairwise squared distance matrix.
    pdist_matrix = metric_learning.pairwise_distance(embeddings, squared=True)
    # Build pairwise binary adjacency matrix.
    adjacency = tf.math.equal(labels, tf.transpose(labels))
    # Invert so we can select negatives only.
    adjacency_not = tf.math.logical_not(adjacency)

    batch_size = tf.size(labels)

    # Compute the mask.
    pdist_matrix_tile = tf.tile(pdist_matrix, [batch_size, 1])
    mask = tf.math.logical_and(
        tf.tile(adjacency_not, [batch_size, 1]),
        tf.math.greater(
            pdist_matrix_tile, tf.reshape(tf.transpose(pdist_matrix), [-1, 1])
        ),
    )
    mask_final = tf.reshape(
        tf.math.greater(
            tf.math.reduce_sum(
                tf.cast(mask, dtype=tf.dtypes.float32), 1, keepdims=True
            ),
            0.0,
        ),
        [batch_size, batch_size],
    )
    mask_final = tf.transpose(mask_final)

    adjacency_not = tf.cast(adjacency_not, dtype=tf.dtypes.float32)
    mask = tf.cast(mask, dtype=tf.dtypes.float32)

    # negatives_outside: smallest D_an where D_an > D_ap.
    negatives_outside = tf.reshape(
        _masked_minimum(pdist_matrix_tile, mask), [batch_size, batch_size]
    )
    negatives_outside = tf.transpose(negatives_outside)

    # negatives_inside: largest D_an.
    negatives_inside = tf.tile(
        _masked_maximum(pdist_matrix, adjacency_not), [1, batch_size]
    )
    semi_hard_negatives = tf.where(mask_final, negatives_outside, negatives_inside)

    loss_mat = tf.math.add(margin, pdist_matrix - semi_hard_negatives)

    mask_positives = tf.cast(adjacency, dtype=tf.dtypes.float32) - tf.linalg.diag(
        tf.ones([batch_size])
    )

    # In lifted-struct, the authors multiply 0.5 for upper triangular
    #   in semihard, they take all positive pairs except the diagonal.
    num_positives = tf.math.reduce_sum(mask_positives)

    triplet_loss = tf.math.truediv(
        tf.math.reduce_sum(
            tf.math.maximum(tf.math.multiply(loss_mat, mask_positives), 0.0)
        ),
        num_positives,
    )

    return triplet_loss


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

# Fixed code
def triplet_semihard_loss(y_true, y_pred, margin=1.0):
    """Computes the triplet loss with semi-hard negative mining.

    Args:
      y_true: 1-D integer `Tensor` with shape [batch_size] of
        multiclass integer labels.
      y_pred: 2-D float `Tensor` of embedding vectors. Embeddings should
        be l2 normalized.
      margin: Float, margin term in the loss definition.
    """
    labels, embeddings = y_true, y_pred
    # Reshape label tensor to [batch_size, 1].
    lshape = tf.shape(labels)
    labels = tf.reshape(labels, [lshape[0], 1])

    # Build pairwise squared distance matrix.
    pdist_matrix = metric_learning.pairwise_distance(embeddings, squared=True)
    # Build pairwise binary adjacency matrix.
    adjacency = tf.math.equal(labels, tf.transpose(labels))
    # Invert so we can select negatives only.
    adjacency_not = tf.math.logical_not(adjacency)

    batch_size = tf.size(labels)

    # Compute the mask.
    pdist_matrix_tile = tf.tile(pdist_matrix, [batch_size, 1])
    mask = tf.math.logical_and(
        tf.tile(adjacency_not, [batch_size, 1]),
        tf.math.greater(
            pdist_matrix_tile, tf.reshape(tf.transpose(pdist_matrix), [-1, 1])
        ),
    )
    mask_final = tf.reshape(
        tf.math.greater(
            tf.math.reduce_sum(
                tf.cast(mask, dtype=tf.dtypes.float32), 1, keepdims=True
            ),
            0.0,
        ),
        [batch_size, batch_size],
    )
    mask_final = tf.transpose(mask_final)

    adjacency_not = tf.cast(adjacency_not, dtype=tf.dtypes.float32)
    mask = tf.cast(mask, dtype=tf.dtypes.float32)

    # negatives_outside: smallest D_an where D_an > D_ap.
    negatives_outside = tf.reshape(
        _masked_minimum(pdist_matrix_tile, mask), [batch_size, batch_size]
    )
    negatives_outside = tf.transpose(negatives_outside)

    # negatives_inside: largest D_an.
    negatives_inside = tf.tile(
        _masked_maximum(pdist_matrix, adjacency_not), [1, batch_size]
    )
    semi_hard_negatives = tf.where(mask_final, negatives_outside, negatives_inside)

    loss_mat = tf.math.add(margin, pdist_matrix - semi_hard_negatives)

    mask_positives = tf.cast(adjacency, dtype=tf.dtypes.float32) - tf.linalg.diag(
        tf.ones([batch_size])
    )

    # In lifted-struct, the authors multiply 0.5 for upper triangular
    #   in semihard, they take all positive pairs except the diagonal.
    num_positives = tf.math.reduce_sum(mask_positives)

    triplet_loss = tf.math.truediv(
        tf.math.reduce_sum(
            tf.math.maximum(tf.math.multiply(loss_mat, mask_positives), 0.0)
        ),
        num_positives,
    )

    return triplet_loss