Repository: tensorflow/addons

Issue:
tfa.seq2seq.sequence_loss can't average over one dimension (batch or timesteps) while summing over the other one
**System information**
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Google Colab
- TensorFlow installed from (source or binary): binary
- TensorFlow version (use command below): 2.0.0=beta1
- TensorFlow Addons installed from (source, PyPi): PyPi
- TensorFlow Addons version: 0.4.0
- Python version and type (eg. Anaconda Python, Stock Python as in Mac, or homebrew installed Python etc): Google Colab Python
- Is GPU used? (yes/no): yes
- GPU model (if used): T4

**Describe the bug**

`tfa.seq2seq.sequence_loss` can't average over one dimension (`batch` or `timesteps`) while summing over the other one. It will arbitrarily only execute the averaging and ignore the sum right now.

**Describe the expected behavior**

I think the weights should be associated with the summing operation, and then the averaging should happen irrespective of that.
Concretely, when passing, say `average_across_batch=True` and `sum_over_timesteps=True` (of course, making sure `average_across_timesteps=False` is set), you should expect either of these things:

1. An error stating that this is not implemented (might be the wisest).
2. Return a scalar tensor obtained by either of these two following orders:
  a) first computing the *weighted sum* of xents over timesteps (yielding a batchsize-sized tensor of xent-sums), then simply averaging this vector, i.e., summing and dividing by the batchsize. The result, however, is just the both-averaged version times the batchsize, divided by the sum of all weights.
  b) first computing the *weighted average* over the batchsize, then summing these averages over all timesteps. The result here is different from 1a and the double-averaged (of course, there is some correlation...)!

I think 1a is the desired behavior (as the loglikelihood of a sequence really is the sum of the individual loglikelihoods and batches do correspond to sequence-length agnostic averages) and I'd be happy to establish it as the standard for this. Either way, doing something other than failing with an error will require an explicit notice in the docs. An error (or warning for backwards-compatibility?) might just be the simplest and safest option.

**Code to reproduce the issue**

```python
tfa.seq2seq.sequence_loss(
  logits=tf.random.normal([3, 5, 7]),
  targets=tf.zeros([3, 5], dtype=tf.int32),
  weights=tf.sequence_mask(lengths=[3, 5, 1], maxlen=5, dtype=tf.float32),
  average_across_batch=True,
  average_across_timesteps=False,
  sum_over_batch=False,
  sum_over_timesteps=True,
)
```
...should return a scalar but returns only the batch-averaged tensor.

**Some more code to play with to test the claims above**

```python
import tensorflow.compat.v2 as tf
import tensorflow_addons as tfa
import numpy as np
import random

case1b = []
dblavg = []

for _ in range(100):
  dtype = tf.float32
  batchsize = random.randint(2, 10)
  maxlen = random.randint(2, 10)
  logits = tf.random.normal([batchsize, maxlen, 3])
  labels = tf.zeros([batchsize, maxlen], dtype=tf.int32)
  lengths = tf.squeeze(tf.random.categorical(tf.zeros([1, maxlen - 1]), batchsize)) + 1
  weights = tf.sequence_mask(lengths=lengths, maxlen=maxlen, dtype=tf.float32)

  def sl(ab, sb, at, st):
    return tfa.seq2seq.sequence_loss(
        logits,
        labels,
        weights,
        average_across_batch=ab,
        average_across_timesteps=at,
        sum_over_batch=sb,
        sum_over_timesteps=st,
    )

  all_b_all_t = sl(ab=False, sb=False, at=False, st=False)
  avg_b_avg_t = sl(ab=True,  sb=False, at=True,  st=False)
  sum_b_all_t = sl(ab=False, sb=True,  at=False, st=False)

  tf.assert_equal(sum_b_all_t, tf.math.divide_no_nan(tf.reduce_sum(all_b_all_t, axis=0), tf.reduce_sum(weights, axis=0)))

  weighted = all_b_all_t * weights

  first_sum_timesteps = tf.reduce_sum(weighted, axis=1)
  then_average_batch = tf.reduce_sum(first_sum_timesteps) / batchsize

  first_average_batch = tf.math.divide_no_nan(tf.reduce_sum(weighted, axis=0), tf.reduce_sum(weights, axis=0))
  then_sum_timesteps = tf.reduce_sum(first_average_batch)

  # Case 1a and 1b are different.
  assert not np.isclose(then_average_batch, then_sum_timesteps)
  # Case 1a is just the double-averaging up to a constant.
  assert     np.allclose(then_average_batch * batchsize / tf.reduce_sum(weights), avg_b_avg_t)
  # Case 1b is not just the averaging.
  assert not np.allclose(then_sum_timesteps / maxlen, avg_b_avg_t)
  # They only kind of correlate:
  case1b.append(then_sum_timesteps / maxlen)
  dblavg.append(avg_b_avg_t)
```


Discussion:
Hi @sjmielke, thanks for the suggestion. I think the error message here makes sense. Feel free to send a PR for this.

Resolving patch:
diff --git a/tensorflow_addons/seq2seq/loss.py b/tensorflow_addons/seq2seq/loss.py
--- a/tensorflow_addons/seq2seq/loss.py
+++ b/tensorflow_addons/seq2seq/loss.py
@@ -103,6 +103,14 @@ def sequence_loss(logits,
         raise ValueError(
             "average_across_batch and sum_over_batch cannot be set "
             "to True at same time.")
+    if average_across_batch and sum_over_timesteps:
+        raise ValueError(
+            "average_across_batch and sum_over_timesteps cannot be set "
+            "to True at same time because of ambiguous order.")
+    if sum_over_batch and average_across_timesteps:
+        raise ValueError(
+            "sum_over_batch and average_across_timesteps cannot be set "
+            "to True at same time because of ambiguous order.")
     with tf.name_scope(name or "sequence_loss"):
         num_classes = tf.shape(input=logits)[2]
         logits_flat = tf.reshape(logits, [-1, num_classes])


Test patch:
diff --git a/tensorflow_addons/seq2seq/loss_test.py b/tensorflow_addons/seq2seq/loss_test.py
--- a/tensorflow_addons/seq2seq/loss_test.py
+++ b/tensorflow_addons/seq2seq/loss_test.py
@@ -297,6 +297,18 @@ def testZeroWeights(self):
             compare_total = np.zeros((self.batch_size, self.sequence_length))
             self.assertAllClose(compare_total, res)
 
+    def testAmbiguousOrder(self):
+        with self.assertRaisesRegexp(ValueError, 'because of ambiguous order'):
+            with self.cached_session(use_gpu=True):
+                self.setup()
+                seq_loss = loss.SequenceLoss(
+                    average_across_timesteps=False,
+                    average_across_batch=True,
+                    sum_over_timesteps=True,
+                    sum_over_batch=False)
+                self.evaluate(
+                    seq_loss(self.targets, self.logits, self.weights))
+
 
 if __name__ == '__main__':
     tf.test.main()