# Buggy code
def pinball_loss(
    y_true: TensorLike, y_pred: TensorLike, tau: FloatTensorLike = 0.5
) -> tf.Tensor:
    """Computes the pinball loss between `y_true` and `y_pred`.

    `loss = maximum(tau * (y_true - y_pred), (tau - 1) * (y_true - y_pred))`

    In the context of regression this, loss yields an estimator of the tau
    conditional quantile.

    See: https://en.wikipedia.org/wiki/Quantile_regression

    Usage:
    ```python
    loss = pinball_loss([0., 0., 1., 1.], [1., 1., 1., 0.], tau=.1)

    # loss = max(0.1 * (y_true - y_pred), (0.1 - 1) * (y_true - y_pred))
    #      = (0.9 + 0.9 + 0 + 0.1) / 4

    print('Loss: ', loss.numpy())  # Loss: 0.475
    ```

    Args:
      y_true: Ground truth values. shape = `[batch_size, d0, .. dN]`
      y_pred: The predicted values. shape = `[batch_size, d0, .. dN]`
      tau: (Optional) Float in [0, 1] or a tensor taking values in [0, 1] and
        shape = `[d0,..., dn]`.  It defines the slope of the pinball loss. In
        the context of quantile regression, the value of tau determines the
        conditional quantile level. When tau = 0.5, this amounts to l1
        regression, an estimator of the conditional median (0.5 quantile).

    Returns:
        pinball_loss: 1-D float `Tensor` with shape [batch_size].

    References:
      - https://en.wikipedia.org/wiki/Quantile_regression
      - https://projecteuclid.org/download/pdfview_1/euclid.bj/1297173840
    """
    y_pred = tf.convert_to_tensor(y_pred)
    y_true = tf.cast(y_true, y_pred.dtype)

    # broadcast the pinball slope along the batch dimension, and clip to
    # acceptable values
    tau = tf.expand_dims(tf.cast(tau, y_pred.dtype), 0)
    one = tf.cast(1, tau.dtype)

    delta_y = y_true - y_pred
    pinball = tf.math.maximum(tau * delta_y, (tau - one) * delta_y)
    return tf.reduce_mean(tf.keras.backend.batch_flatten(pinball), axis=-1)


# Traceback (tensorflow.python.framework.errors_impl.InvalidArgumentError)
(proof_of_concept_pytest) /mnt/c/Users/gdemarmi/Desktop/projects/addons $ pytest -v tensorflow_addons/losses/quantiles_test.py
/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/pep8.py:110: FutureWarning: Possible nested set at position 1
EXTRANEOUS_WHITESPACE_REGEX = re.compile(r'[[({] | []}),;:]')
================================================================= test session starts ==================================================================
platform linux -- Python 3.7.5, pytest-5.0.1, py-1.8.0, pluggy-0.13.0 -- /home/gdemarmi/softwares/python/anaconda/bin/python
cachedir: .pytest_cache
rootdir: /mnt/c/Users/gdemarmi/Desktop/projects/addons
plugins: arraydiff-0.3, cov-2.7.1, doctestplus-0.2.0, flake8-1.0.4, forked-1.0.2, openfiles-0.4.0, pep8-1.0.6, remotedata-0.3.2, xdist-1.27.0, typeguard-2.7.1
collected 15 items

tensorflow_addons/losses/quantiles_test.py::PinballLossTest::test_all_correct_unweighted PASSED                                                  [  6%]
tensorflow_addons/losses/quantiles_test.py::PinballLossTest::test_config PASSED                                                                  [ 13%]
tensorflow_addons/losses/quantiles_test.py::PinballLossTest::test_invalid_sample_weight FAILED                                                   [ 20%]
tensorflow_addons/losses/quantiles_test.py::PinballLossTest::test_no_reduction PASSED                                                            [ 26%]
tensorflow_addons/losses/quantiles_test.py::PinballLossTest::test_sample_weighted PASSED                                                         [ 33%]
tensorflow_addons/losses/quantiles_test.py::PinballLossTest::test_scalar_weighted PASSED                                                         [ 40%]
tensorflow_addons/losses/quantiles_test.py::PinballLossTest::test_session SKIPPED                                                                [ 46%]
tensorflow_addons/losses/quantiles_test.py::PinballLossTest::test_sum_reduction PASSED                                                           [ 53%]
tensorflow_addons/losses/quantiles_test.py::PinballLossTest::test_timestep_weighted FAILED                                                       [ 60%]
tensorflow_addons/losses/quantiles_test.py::PinballLossTest::test_unweighted PASSED                                                              [ 66%]
tensorflow_addons/losses/quantiles_test.py::PinballLossTest::test_unweighted_quantile_0pc PASSED                                                 [ 73%]
tensorflow_addons/losses/quantiles_test.py::PinballLossTest::test_unweighted_quantile_100pc PASSED                                               [ 80%]
tensorflow_addons/losses/quantiles_test.py::PinballLossTest::test_unweighted_quantile_10pc PASSED                                                [ 86%]
tensorflow_addons/losses/quantiles_test.py::PinballLossTest::test_unweighted_quantile_90pc PASSED                                                [ 93%]
tensorflow_addons/losses/quantiles_test.py::PinballLossTest::test_zero_weighted PASSED                                                           [100%]

======================================================================= FAILURES =======================================================================
______________________________________________________ PinballLossTest.test_invalid_sample_weight ______________________________________________________

def _create_c_op(graph, node_def, inputs, control_inputs):
"""Creates a TF_Operation.

Args:
graph: a `Graph`.
node_def: `node_def_pb2.NodeDef` for the operation to create.
inputs: A list of `Tensor`s (corresponding to scalar inputs) and lists of
`Tensor`s (corresponding to sequence inputs, e.g. "int64 * N",
"list(int64)"). The length of the list should be equal to the number of
inputs specified by this operation's op def.
control_inputs: A list of `Operation`s to set as control dependencies.

Returns:
A wrapped TF_Operation*.
"""
# pylint: disable=protected-access
op_desc = c_api.TF_NewOperation(graph._c_graph, compat.as_str(node_def.op),
compat.as_str(node_def.name))
if node_def.device:
c_api.TF_SetDevice(op_desc, compat.as_str(node_def.device))
# Add inputs
for op_input in inputs:
if isinstance(op_input, (list, tuple)):
c_api.TF_AddInputList(op_desc, [t._as_tf_output() for t in op_input])
else:
c_api.TF_AddInput(op_desc, op_input._as_tf_output())

# Add control inputs
for control_input in control_inputs:
c_api.TF_AddControlInput(op_desc, control_input._c_op)
# pylint: enable=protected-access

# Add attrs
for name, attr_value in node_def.attr.items():
serialized = attr_value.SerializeToString()
# TODO(skyewm): this creates and deletes a new TF_Status for every attr.
# It might be worth creating a convenient way to re-use the same status.
c_api.TF_SetAttrValueProto(op_desc, compat.as_str(name), serialized)

try:
c_op = c_api.TF_FinishOperation(op_desc)
E       tensorflow.python.framework.errors_impl.InvalidArgumentError: Can not squeeze dim[1], expected a dimension of 1, got 2 for 'pinball_loss/weighted_loss/Squeeze' (op: 'Squeeze') with input shapes: [2,2].

/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py:1619: InvalidArgumentError

During handling of the above exception, another exception occurred:
Traceback (most recent call last):
File "/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py", line 1619, in _create_c_op
c_op = c_api.TF_FinishOperation(op_desc)
tensorflow.python.framework.errors_impl.InvalidArgumentError: Can not squeeze dim[1], expected a dimension of 1, got 2 for 'pinball_loss/weighted_loss/Squeeze' (op: 'Squeeze') with input shapes: [2,2].

During handling of the above exception, another exception occurred:

ValueError: Can not squeeze dim[1], expected a dimension of 1, got 2 for 'pinball_loss/weighted_loss/Squeeze' (op: 'Squeeze') with input shapes: [2,2].

During handling of the above exception, another exception occurred:

self = <tensorflow_addons.losses.quantiles_test.PinballLossTest testMethod=test_invalid_sample_weight>, args = (), kwargs = {}

def decorated(self, *args, **kwargs):
try:
with context.graph_mode():
with self.test_session(use_gpu=use_gpu, config=config):
f(self, *args, **kwargs)

/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/tensorflow_core/python/framework/test_util.py:1111:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
tensorflow_addons/losses/quantiles_test.py:113: in test_invalid_sample_weight
pin_obj(y_true, y_pred, sample_weight=sample_weight)
E   AssertionError: "weights can not be broadcast to values" does not match "Can not squeeze dim[1], expected a dimension of 1, got 2 for 'pinball_loss/weighted_loss/Squeeze' (op: 'Squeeze') with input shapes: [2,2]."
________________________________________________________ PinballLossTest.test_timestep_weighted ________________________________________________________

graph = <tensorflow.python.framework.ops.Graph object at 0x7f562c26e950>
node_def = name: "pinball_loss/weighted_loss/Squeeze"
op: "Squeeze"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "squeeze_dims"
value {
list {
i: -1
}
}
}

inputs = [<tf.Tensor 'pinball_loss/weighted_loss/Cast:0' shape=(2, 3) dtype=float32>], control_inputs = []

def _create_c_op(graph, node_def, inputs, control_inputs):
"""Creates a TF_Operation.

Args:
graph: a `Graph`.
node_def: `node_def_pb2.NodeDef` for the operation to create.
inputs: A list of `Tensor`s (corresponding to scalar inputs) and lists of
`Tensor`s (corresponding to sequence inputs, e.g. "int64 * N",
"list(int64)"). The length of the list should be equal to the number of
inputs specified by this operation's op def.
control_inputs: A list of `Operation`s to set as control dependencies.

Returns:
A wrapped TF_Operation*.
"""
# pylint: disable=protected-access
op_desc = c_api.TF_NewOperation(graph._c_graph, compat.as_str(node_def.op),
compat.as_str(node_def.name))
if node_def.device:
c_api.TF_SetDevice(op_desc, compat.as_str(node_def.device))
# Add inputs
for op_input in inputs:
if isinstance(op_input, (list, tuple)):
c_api.TF_AddInputList(op_desc, [t._as_tf_output() for t in op_input])
else:
c_api.TF_AddInput(op_desc, op_input._as_tf_output())

# Add control inputs
for control_input in control_inputs:
c_api.TF_AddControlInput(op_desc, control_input._c_op)
# pylint: enable=protected-access

# Add attrs
for name, attr_value in node_def.attr.items():
serialized = attr_value.SerializeToString()
# TODO(skyewm): this creates and deletes a new TF_Status for every attr.
# It might be worth creating a convenient way to re-use the same status.
c_api.TF_SetAttrValueProto(op_desc, compat.as_str(name), serialized)

try:
c_op = c_api.TF_FinishOperation(op_desc)
E       tensorflow.python.framework.errors_impl.InvalidArgumentError: Can not squeeze dim[1], expected a dimension of 1, got 3 for 'pinball_loss/weighted_loss/Squeeze' (op: 'Squeeze') with input shapes: [2,3].

/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py:1619: InvalidArgumentError

During handling of the above exception, another exception occurred:

self = <tensorflow_addons.losses.quantiles_test.PinballLossTest testMethod=test_timestep_weighted>, args = (), kwargs = {}

def decorated(self, *args, **kwargs):
try:
with context.graph_mode():
with self.test_session(use_gpu=use_gpu, config=config):
f(self, *args, **kwargs)

/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/tensorflow_core/python/framework/test_util.py:1111:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
tensorflow_addons/losses/quantiles_test.py:95: in test_timestep_weighted
loss = pin_obj(y_true, y_pred, sample_weight=sample_weight)
/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/tensorflow_core/python/keras/losses.py:128: in __call__
losses, sample_weight, reduction=self._get_reduction())
/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/tensorflow_core/python/keras/utils/losses_utils.py:107: in compute_weighted_loss
losses, sample_weight)
/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/tensorflow_core/python/ops/losses/util.py:145: in scale_losses_by_sample_weight
losses, None, sample_weight)
/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/tensorflow_core/python/ops/losses/util.py:97: in squeeze_or_expand_dimensions
sample_weight = array_ops.squeeze(sample_weight, [-1])
/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/tensorflow_core/python/util/dispatch.py:180: in wrapper
return target(*args, **kwargs)
/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/tensorflow_core/python/util/deprecation.py:507: in new_func
return func(*args, **kwargs)
/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/tensorflow_core/python/ops/array_ops.py:3780: in squeeze
return gen_array_ops.squeeze(input, axis, name)
/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/tensorflow_core/python/ops/gen_array_ops.py:9231: in squeeze
"Squeeze", input=input, squeeze_dims=axis, name=name)
/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/tensorflow_core/python/framework/op_def_library.py:742: in _apply_op_helper
attrs=attr_protos, op_def=op_def)
/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py:3322: in _create_op_internal
op_def=op_def)
/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py:1786: in __init__
control_input_ops)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

graph = <tensorflow.python.framework.ops.Graph object at 0x7f562c26e950>
node_def = name: "pinball_loss/weighted_loss/Squeeze"
op: "Squeeze"
attr {
key: "T"
value {
type: DT_FLOAT
}
}
attr {
key: "squeeze_dims"
value {
list {
i: -1
}
}
}

inputs = [<tf.Tensor 'pinball_loss/weighted_loss/Cast:0' shape=(2, 3) dtype=float32>], control_inputs = []

def _create_c_op(graph, node_def, inputs, control_inputs):
"""Creates a TF_Operation.

Args:
graph: a `Graph`.
node_def: `node_def_pb2.NodeDef` for the operation to create.
inputs: A list of `Tensor`s (corresponding to scalar inputs) and lists of
`Tensor`s (corresponding to sequence inputs, e.g. "int64 * N",
"list(int64)"). The length of the list should be equal to the number of
inputs specified by this operation's op def.
control_inputs: A list of `Operation`s to set as control dependencies.

Returns:
A wrapped TF_Operation*.
"""
# pylint: disable=protected-access
op_desc = c_api.TF_NewOperation(graph._c_graph, compat.as_str(node_def.op),
compat.as_str(node_def.name))
if node_def.device:
c_api.TF_SetDevice(op_desc, compat.as_str(node_def.device))
# Add inputs
for op_input in inputs:
if isinstance(op_input, (list, tuple)):
c_api.TF_AddInputList(op_desc, [t._as_tf_output() for t in op_input])
else:
c_api.TF_AddInput(op_desc, op_input._as_tf_output())

# Add control inputs
for control_input in control_inputs:
c_api.TF_AddControlInput(op_desc, control_input._c_op)
# pylint: enable=protected-access

# Add attrs
for name, attr_value in node_def.attr.items():
serialized = attr_value.SerializeToString()
# TODO(skyewm): this creates and deletes a new TF_Status for every attr.
# It might be worth creating a convenient way to re-use the same status.
c_api.TF_SetAttrValueProto(op_desc, compat.as_str(name), serialized)

try:
c_op = c_api.TF_FinishOperation(op_desc)
except errors.InvalidArgumentError as e:
# Convert to ValueError for backwards compatibility.
raise ValueError(str(e))
E       ValueError: Can not squeeze dim[1], expected a dimension of 1, got 3 for 'pinball_loss/weighted_loss/Squeeze' (op: 'Squeeze') with input shapes: [2,3].

/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/tensorflow_core/python/framework/ops.py:1622: ValueError
----------------------------------------------------------------- Captured stderr call -----------------------------------------------------------------
W0302 15:37:30.724358 140009831860032 def_function.py:586] 8 out of the last 11 calls to <function pinball_loss at 0x7f562f7ecb00> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings is likely due to passing python objects instead of tensors. Also, tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. Please refer to https://www.tensorflow.org/tutorials/customization/performance#python_or_tensor_args and https://www.tensorflow.org/api_docs/python/tf/function for more details.
------------------------------------------------------------------ Captured log call -------------------------------------------------------------------
WARNING  tensorflow:def_function.py:586 8 out of the last 11 calls to <function pinball_loss at 0x7f562f7ecb00> triggered tf.function retracing. Tracing is expensive and the excessive number of tracings is likely due to passing python objects instead of tensors. Also, tf.function has experimental_relax_shapes=True option that relaxes argument shapes that can avoid unnecessary retracing. Please refer to https://www.tensorflow.org/tutorials/customization/performance#python_or_tensor_args and https://www.tensorflow.org/api_docs/python/tf/function for more details.
=================================================================== warnings summary ===================================================================
/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/google/protobuf/descriptor.py:47
/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/google/protobuf/descriptor.py:47: DeprecationWarning: Using or importing the ABCs from 'collections' instead of from 'collections.abc' is deprecated since Python 3.3,and in 3.9 it will stop working
from google.protobuf.pyext import _message

/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/botocore/vendored/requests/packages/urllib3/_collections.py:1
/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/botocore/vendored/requests/packages/urllib3/_collections.py:1
/home/gdemarmi/softwares/python/anaconda/lib/python3.7/site-packages/botocore/vendored/requests/packages/urllib3/_collections.py:1: DeprecationWarning: Using or importing the ABCs from 'collections' instead of from 'collections.abc' is deprecated since Python 3.3,and in 3.9 it will stop working
from collections import Mapping, MutableMapping

-- Docs: https://docs.pytest.org/en/latest/warnings.html
============================================= 2 failed, 12 passed, 1 skipped, 3 warnings in 17.15 seconds ==============================================

# Fixed code
def pinball_loss(
    y_true: TensorLike, y_pred: TensorLike, tau: FloatTensorLike = 0.5
) -> tf.Tensor:
    """Computes the pinball loss between `y_true` and `y_pred`.

    `loss = maximum(tau * (y_true - y_pred), (tau - 1) * (y_true - y_pred))`

    In the context of regression this, loss yields an estimator of the tau
    conditional quantile.

    See: https://en.wikipedia.org/wiki/Quantile_regression

    Usage:
    ```python
    loss = pinball_loss([0., 0., 1., 1.], [1., 1., 1., 0.], tau=.1)

    # loss = max(0.1 * (y_true - y_pred), (0.1 - 1) * (y_true - y_pred))
    #      = (0.9 + 0.9 + 0 + 0.1) / 4

    print('Loss: ', loss.numpy())  # Loss: 0.475
    ```

    Args:
      y_true: Ground truth values. shape = `[batch_size, d0, .. dN]`
      y_pred: The predicted values. shape = `[batch_size, d0, .. dN]`
      tau: (Optional) Float in [0, 1] or a tensor taking values in [0, 1] and
        shape = `[d0,..., dn]`.  It defines the slope of the pinball loss. In
        the context of quantile regression, the value of tau determines the
        conditional quantile level. When tau = 0.5, this amounts to l1
        regression, an estimator of the conditional median (0.5 quantile).

    Returns:
        pinball_loss: 1-D float `Tensor` with shape [batch_size].

    References:
      - https://en.wikipedia.org/wiki/Quantile_regression
      - https://projecteuclid.org/download/pdfview_1/euclid.bj/1297173840
    """
    y_pred = tf.convert_to_tensor(y_pred)
    y_true = tf.cast(y_true, y_pred.dtype)

    # Broadcast the pinball slope along the batch dimension
    tau = tf.expand_dims(tf.cast(tau, y_pred.dtype), 0)
    one = tf.cast(1, tau.dtype)

    delta_y = y_true - y_pred
    pinball = tf.math.maximum(tau * delta_y, (tau - one) * delta_y)
    return tf.reduce_mean(pinball, axis=-1)