Repository: tensorflow/addons

Issue:
New SelectV2 is causing nightly failures
Since https://github.com/tensorflow/tensorflow/commit/01207a94c627a1bc02e85ab260fcd16c78a0e8da we're seeing failures across a few modules for things that are not broadcastable. 

I will dive in a little deeper, but I was wondering @yongtang if you could advise us to see if there is something obvious that needs to be changed?

https://source.cloud.google.com/results/invocations/3e5ada79-f40b-4033-be80-2f52a382bd5d/targets/tensorflow_addons%2Fubuntu%2Fcpu%2Fpy2%2Fcontinuous/log

```
  File "/tmpfs/src/github/tensorflow_addons/tf/local/lib/python2.7/site-packages/tensorflow/python/ops/array_ops.py", line 3312, in where_v2
    return gen_math_ops.select_v2(condition=condition, t=x, e=y, name=name)
  File "/tmpfs/src/github/tensorflow_addons/tf/local/lib/python2.7/site-packages/tensorflow/python/ops/gen_math_ops.py", line 9290, in select_v2
    _six.raise_from(_core._status_to_exception(e.code, message), None)
  File "/tmpfs/src/github/tensorflow_addons/tf/local/lib/python2.7/site-packages/six.py", line 737, in raise_from
    raise value
InvalidArgumentError: condition [5], then [5,7], and else [5,7] must be broadcastable [Op:SelectV2]
```


Discussion:
@seanpmorgan The new where follows the same broadcast rule as numpy (and other ops in tensorflow):
```
>>> c = np.random.rand(5)
>>> t = np.random.rand(5, 7)
>>> e = np.random.rand(5, 7)
>>> np.where(c, t, e)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ValueError: shape mismatch: objects cannot be broadcast to a single shape
```

The following will work:
```
>>> c = np.random.rand(7)
>>> t = np.random.rand(5, 7)
>>> e = np.random.rand(5, 7)
>>> np.where(c, t, e)
array([[ 0.04187991,  0.26753897,  0.59514661,  0.79760828,  0.88163072,
         0.76875137,  0.87185519],
       [ 0.49383641,  0.9541652 ,  0.65635514,  0.90680339,  0.01328923,
         0.29056175,  0.60653385],
       [ 0.04124529,  0.2395378 ,  0.36566755,  0.27227543,  0.20361383,
         0.5402641 ,  0.39675447],
       [ 0.9076939 ,  0.66453692,  0.23566607,  0.7891645 ,  0.71989783,
         0.45527365,  0.36589332],
       [ 0.88392814,  0.587176  ,  0.0532821 ,  0.19238287,  0.24158712,
         0.60021316,  0.95759263]])
>>> 
```

Resolving patch:
diff --git a/tensorflow_addons/activations/sparsemax.py b/tensorflow_addons/activations/sparsemax.py
--- a/tensorflow_addons/activations/sparsemax.py
+++ b/tensorflow_addons/activations/sparsemax.py
@@ -126,7 +126,8 @@ def _compute_2d_sparsemax(logits, name=None):
     p = tf.math.maximum(
         tf.cast(0, logits.dtype), z - tf.expand_dims(tau_z, -1))
     # If k_z = 0 or if z = nan, then the input is invalid
-    p_safe = tf.where(
+    # TODO: Adjust dimension order for TF2 broadcasting
+    p_safe = tf.compat.v1.where(
         tf.math.logical_or(
             tf.math.equal(k_z, 0), tf.math.is_nan(z_cumsum[:, -1])),
         tf.fill([obs, dims], tf.cast(float("nan"), logits.dtype)), p)
diff --git a/tensorflow_addons/seq2seq/decoder.py b/tensorflow_addons/seq2seq/decoder.py
--- a/tensorflow_addons/seq2seq/decoder.py
+++ b/tensorflow_addons/seq2seq/decoder.py
@@ -406,8 +406,10 @@ def body(time, outputs_ta, state, inputs, finished, sequence_lengths):
             # Zero out output values past finish
             if impute_finished:
                 emit = tf.nest.map_structure(
-                    lambda out, zero: tf.where(finished, zero, out),
-                    next_outputs, zero_outputs)
+                    # TODO: Adjust dimension order for TF2 broadcasting
+                    lambda out, zero: tf.compat.v1.where(finished, zero, out),
+                    next_outputs,
+                    zero_outputs)
             else:
                 emit = next_outputs
 
@@ -419,7 +421,9 @@ def _maybe_copy_state(new, cur):
                 else:
                     new.set_shape(cur.shape)
                     pass_through = (new.shape.ndims == 0)
-                return new if pass_through else tf.where(finished, cur, new)
+                # TODO: Adjust dimension order for TF2 broadcasting
+                return new if pass_through else tf.compat.v1.where(
+                    finished, cur, new)
 
             if impute_finished:
                 next_state = tf.nest.map_structure(_maybe_copy_state,
diff --git a/tensorflow_addons/seq2seq/sampler.py b/tensorflow_addons/seq2seq/sampler.py
--- a/tensorflow_addons/seq2seq/sampler.py
+++ b/tensorflow_addons/seq2seq/sampler.py
@@ -448,9 +448,10 @@ def maybe_concatenate_auxiliary_inputs(outputs_, indices=None):
                     auxiliary_inputs)
 
             if self.next_inputs_fn is None:
-                return tf.where(sample_ids,
-                                maybe_concatenate_auxiliary_inputs(outputs),
-                                base_next_inputs)
+                # TODO: Adjust dimension order for TF2 broadcasting
+                return tf.compat.v1.where(
+                    sample_ids, maybe_concatenate_auxiliary_inputs(outputs),
+                    base_next_inputs)
 
             where_sampling = tf.cast(tf.where(sample_ids), tf.int32)
             where_not_sampling = tf.cast(


Test patch: