We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Update sparsemax to use tf.where V2
As described in #250 and temporarily patched in #251 sparsemax has one instance of tf.where that needs the broadcasting dimensions changed to match numpy and TF2 style.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `tensorflow_addons/activations/sparsemax.py`
Content:
```
1 # Copyright 2016 The TensorFlow Authors. All Rights Reserved.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 # ==============================================================================
15 
16 from __future__ import absolute_import
17 from __future__ import division
18 from __future__ import print_function
19 
20 import tensorflow as tf
21 
22 from tensorflow_addons.utils import keras_utils
23 
24 
25 @tf.function
26 @keras_utils.register_keras_custom_object
27 def sparsemax(logits, axis=-1, name=None):
28     """Sparsemax activation function [1].
29 
30     For each batch `i` and class `j` we have
31       $$sparsemax[i, j] = max(logits[i, j] - tau(logits[i, :]), 0)$$
32 
33     [1]: https://arxiv.org/abs/1602.02068
34 
35     Args:
36         logits: Input tensor.
37         axis: Integer, axis along which the sparsemax operation is applied.
38         name: A name for the operation (optional).
39     Returns:
40         Tensor, output of sparsemax transformation. Has the same type and
41         shape as `logits`.
42     Raises:
43         ValueError: In case `dim(logits) == 1`.
44     """
45     logits = tf.convert_to_tensor(logits, name="logits")
46 
47     # We need its original shape for shape inference.
48     shape = logits.get_shape()
49     rank = shape.rank
50     is_last_axis = (axis == -1) or (axis == rank - 1)
51 
52     if is_last_axis:
53         output = _compute_2d_sparsemax(logits, name=name)
54         output.set_shape(shape)
55         return output
56 
57     # If dim is not the last dimension, we have to do a transpose so that we can
58     # still perform softmax on its last dimension.
59 
60     # Swap logits' dimension of dim and its last dimension.
61     rank_op = tf.rank(logits)
62     axis_norm = axis % rank
63     logits = _swap_axis(logits, axis_norm, tf.math.subtract(rank_op, 1))
64 
65     # Do the actual softmax on its last dimension.
66     output = _compute_2d_sparsemax(logits)
67     output = _swap_axis(
68         output, axis_norm, tf.math.subtract(rank_op, 1), name=name)
69 
70     # Make shape inference work since transpose may erase its static shape.
71     output.set_shape(shape)
72     return output
73 
74 
75 def _swap_axis(logits, dim_index, last_index, **kwargs):
76     return tf.transpose(
77         logits,
78         tf.concat([
79             tf.range(dim_index), [last_index],
80             tf.range(dim_index + 1, last_index), [dim_index]
81         ], 0), **kwargs)
82 
83 
84 @tf.function
85 def _compute_2d_sparsemax(logits, name=None):
86     """Performs the sparsemax operation when axis=-1."""
87     shape_op = tf.shape(logits)
88     obs = tf.math.reduce_prod(shape_op[:-1])
89     dims = shape_op[-1]
90 
91     # In the paper, they call the logits z.
92     # The mean(logits) can be substracted from logits to make the algorithm
93     # more numerically stable. the instability in this algorithm comes mostly
94     # from the z_cumsum. Substacting the mean will cause z_cumsum to be close
95     # to zero. However, in practise the numerical instability issues are very
96     # minor and substacting the mean causes extra issues with inf and nan
97     # input.
98     # Reshape to [obs, dims] as it is almost free and means the remanining
99     # code doesn't need to worry about the rank.
100     z = tf.reshape(logits, [obs, dims])
101 
102     # sort z
103     z_sorted, _ = tf.nn.top_k(z, k=dims)
104 
105     # calculate k(z)
106     z_cumsum = tf.math.cumsum(z_sorted, axis=-1)
107     k = tf.range(1, tf.cast(dims, logits.dtype) + 1, dtype=logits.dtype)
108     z_check = 1 + k * z_sorted > z_cumsum
109     # because the z_check vector is always [1,1,...1,0,0,...0] finding the
110     # (index + 1) of the last `1` is the same as just summing the number of 1.
111     k_z = tf.math.reduce_sum(tf.cast(z_check, tf.int32), axis=-1)
112 
113     # calculate tau(z)
114     # If there are inf values or all values are -inf, the k_z will be zero,
115     # this is mathematically invalid and will also cause the gather_nd to fail.
116     # Prevent this issue for now by setting k_z = 1 if k_z = 0, this is then
117     # fixed later (see p_safe) by returning p = nan. This results in the same
118     # behavior as softmax.
119     k_z_safe = tf.math.maximum(k_z, 1)
120     indices = tf.stack(
121         [tf.range(0, obs), tf.reshape(k_z_safe, [-1]) - 1], axis=1)
122     tau_sum = tf.gather_nd(z_cumsum, indices)
123     tau_z = (tau_sum - 1) / tf.cast(k_z, logits.dtype)
124 
125     # calculate p
126     p = tf.math.maximum(
127         tf.cast(0, logits.dtype), z - tf.expand_dims(tau_z, -1))
128     # If k_z = 0 or if z = nan, then the input is invalid
129     # TODO: Adjust dimension order for TF2 broadcasting
130     p_safe = tf.compat.v1.where(
131         tf.math.logical_or(
132             tf.math.equal(k_z, 0), tf.math.is_nan(z_cumsum[:, -1])),
133         tf.fill([obs, dims], tf.cast(float("nan"), logits.dtype)), p)
134 
135     # Reshape back to original size
136     p_safe = tf.reshape(p_safe, shape_op, name=name)
137     return p_safe
138
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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,11 +126,12 @@
     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
-    # 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)
+    p_safe = tf.where(
+        tf.expand_dims(
+            tf.math.logical_or(
+                tf.math.equal(k_z, 0), tf.math.is_nan(z_cumsum[:, -1])),
+            axis=-1), tf.fill([obs, dims], tf.cast(float("nan"),
+                                                   logits.dtype)), p)
 
     # Reshape back to original size
     p_safe = tf.reshape(p_safe, shape_op, name=name)

{"golden_diff": "diff --git a/tensorflow_addons/activations/sparsemax.py b/tensorflow_addons/activations/sparsemax.py\n--- a/tensorflow_addons/activations/sparsemax.py\n+++ b/tensorflow_addons/activations/sparsemax.py\n@@ -126,11 +126,12 @@\n     p = tf.math.maximum(\n         tf.cast(0, logits.dtype), z - tf.expand_dims(tau_z, -1))\n     # If k_z = 0 or if z = nan, then the input is invalid\n-    # TODO: Adjust dimension order for TF2 broadcasting\n-    p_safe = tf.compat.v1.where(\n-        tf.math.logical_or(\n-            tf.math.equal(k_z, 0), tf.math.is_nan(z_cumsum[:, -1])),\n-        tf.fill([obs, dims], tf.cast(float(\"nan\"), logits.dtype)), p)\n+    p_safe = tf.where(\n+        tf.expand_dims(\n+            tf.math.logical_or(\n+                tf.math.equal(k_z, 0), tf.math.is_nan(z_cumsum[:, -1])),\n+            axis=-1), tf.fill([obs, dims], tf.cast(float(\"nan\"),\n+                                                   logits.dtype)), p)\n \n     # Reshape back to original size\n     p_safe = tf.reshape(p_safe, shape_op, name=name)\n", "issue": "Update sparsemax to use tf.where V2\nAs described in #250 and temporarily patched in #251 sparsemax has one instance of tf.where that needs the broadcasting dimensions changed to match numpy and TF2 style.\n", "before_files": [{"content": "# Copyright 2016 The TensorFlow Authors. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ==============================================================================\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport tensorflow as tf\n\nfrom tensorflow_addons.utils import keras_utils\n\n\n@tf.function\n@keras_utils.register_keras_custom_object\ndef sparsemax(logits, axis=-1, name=None):\n    \"\"\"Sparsemax activation function [1].\n\n    For each batch `i` and class `j` we have\n      $$sparsemax[i, j] = max(logits[i, j] - tau(logits[i, :]), 0)$$\n\n    [1]: https://arxiv.org/abs/1602.02068\n\n    Args:\n        logits: Input tensor.\n        axis: Integer, axis along which the sparsemax operation is applied.\n        name: A name for the operation (optional).\n    Returns:\n        Tensor, output of sparsemax transformation. Has the same type and\n        shape as `logits`.\n    Raises:\n        ValueError: In case `dim(logits) == 1`.\n    \"\"\"\n    logits = tf.convert_to_tensor(logits, name=\"logits\")\n\n    # We need its original shape for shape inference.\n    shape = logits.get_shape()\n    rank = shape.rank\n    is_last_axis = (axis == -1) or (axis == rank - 1)\n\n    if is_last_axis:\n        output = _compute_2d_sparsemax(logits, name=name)\n        output.set_shape(shape)\n        return output\n\n    # If dim is not the last dimension, we have to do a transpose so that we can\n    # still perform softmax on its last dimension.\n\n    # Swap logits' dimension of dim and its last dimension.\n    rank_op = tf.rank(logits)\n    axis_norm = axis % rank\n    logits = _swap_axis(logits, axis_norm, tf.math.subtract(rank_op, 1))\n\n    # Do the actual softmax on its last dimension.\n    output = _compute_2d_sparsemax(logits)\n    output = _swap_axis(\n        output, axis_norm, tf.math.subtract(rank_op, 1), name=name)\n\n    # Make shape inference work since transpose may erase its static shape.\n    output.set_shape(shape)\n    return output\n\n\ndef _swap_axis(logits, dim_index, last_index, **kwargs):\n    return tf.transpose(\n        logits,\n        tf.concat([\n            tf.range(dim_index), [last_index],\n            tf.range(dim_index + 1, last_index), [dim_index]\n        ], 0), **kwargs)\n\n\n@tf.function\ndef _compute_2d_sparsemax(logits, name=None):\n    \"\"\"Performs the sparsemax operation when axis=-1.\"\"\"\n    shape_op = tf.shape(logits)\n    obs = tf.math.reduce_prod(shape_op[:-1])\n    dims = shape_op[-1]\n\n    # In the paper, they call the logits z.\n    # The mean(logits) can be substracted from logits to make the algorithm\n    # more numerically stable. the instability in this algorithm comes mostly\n    # from the z_cumsum. Substacting the mean will cause z_cumsum to be close\n    # to zero. However, in practise the numerical instability issues are very\n    # minor and substacting the mean causes extra issues with inf and nan\n    # input.\n    # Reshape to [obs, dims] as it is almost free and means the remanining\n    # code doesn't need to worry about the rank.\n    z = tf.reshape(logits, [obs, dims])\n\n    # sort z\n    z_sorted, _ = tf.nn.top_k(z, k=dims)\n\n    # calculate k(z)\n    z_cumsum = tf.math.cumsum(z_sorted, axis=-1)\n    k = tf.range(1, tf.cast(dims, logits.dtype) + 1, dtype=logits.dtype)\n    z_check = 1 + k * z_sorted > z_cumsum\n    # because the z_check vector is always [1,1,...1,0,0,...0] finding the\n    # (index + 1) of the last `1` is the same as just summing the number of 1.\n    k_z = tf.math.reduce_sum(tf.cast(z_check, tf.int32), axis=-1)\n\n    # calculate tau(z)\n    # If there are inf values or all values are -inf, the k_z will be zero,\n    # this is mathematically invalid and will also cause the gather_nd to fail.\n    # Prevent this issue for now by setting k_z = 1 if k_z = 0, this is then\n    # fixed later (see p_safe) by returning p = nan. This results in the same\n    # behavior as softmax.\n    k_z_safe = tf.math.maximum(k_z, 1)\n    indices = tf.stack(\n        [tf.range(0, obs), tf.reshape(k_z_safe, [-1]) - 1], axis=1)\n    tau_sum = tf.gather_nd(z_cumsum, indices)\n    tau_z = (tau_sum - 1) / tf.cast(k_z, logits.dtype)\n\n    # calculate p\n    p = tf.math.maximum(\n        tf.cast(0, logits.dtype), z - tf.expand_dims(tau_z, -1))\n    # If k_z = 0 or if z = nan, then the input is invalid\n    # TODO: Adjust dimension order for TF2 broadcasting\n    p_safe = tf.compat.v1.where(\n        tf.math.logical_or(\n            tf.math.equal(k_z, 0), tf.math.is_nan(z_cumsum[:, -1])),\n        tf.fill([obs, dims], tf.cast(float(\"nan\"), logits.dtype)), p)\n\n    # Reshape back to original size\n    p_safe = tf.reshape(p_safe, shape_op, name=name)\n    return p_safe\n", "path": "tensorflow_addons/activations/sparsemax.py"}], "after_files": [{"content": "# Copyright 2016 The TensorFlow Authors. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ==============================================================================\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport tensorflow as tf\n\nfrom tensorflow_addons.utils import keras_utils\n\n\n@tf.function\n@keras_utils.register_keras_custom_object\ndef sparsemax(logits, axis=-1, name=None):\n    \"\"\"Sparsemax activation function [1].\n\n    For each batch `i` and class `j` we have\n      $$sparsemax[i, j] = max(logits[i, j] - tau(logits[i, :]), 0)$$\n\n    [1]: https://arxiv.org/abs/1602.02068\n\n    Args:\n        logits: Input tensor.\n        axis: Integer, axis along which the sparsemax operation is applied.\n        name: A name for the operation (optional).\n    Returns:\n        Tensor, output of sparsemax transformation. Has the same type and\n        shape as `logits`.\n    Raises:\n        ValueError: In case `dim(logits) == 1`.\n    \"\"\"\n    logits = tf.convert_to_tensor(logits, name=\"logits\")\n\n    # We need its original shape for shape inference.\n    shape = logits.get_shape()\n    rank = shape.rank\n    is_last_axis = (axis == -1) or (axis == rank - 1)\n\n    if is_last_axis:\n        output = _compute_2d_sparsemax(logits, name=name)\n        output.set_shape(shape)\n        return output\n\n    # If dim is not the last dimension, we have to do a transpose so that we can\n    # still perform softmax on its last dimension.\n\n    # Swap logits' dimension of dim and its last dimension.\n    rank_op = tf.rank(logits)\n    axis_norm = axis % rank\n    logits = _swap_axis(logits, axis_norm, tf.math.subtract(rank_op, 1))\n\n    # Do the actual softmax on its last dimension.\n    output = _compute_2d_sparsemax(logits)\n    output = _swap_axis(\n        output, axis_norm, tf.math.subtract(rank_op, 1), name=name)\n\n    # Make shape inference work since transpose may erase its static shape.\n    output.set_shape(shape)\n    return output\n\n\ndef _swap_axis(logits, dim_index, last_index, **kwargs):\n    return tf.transpose(\n        logits,\n        tf.concat([\n            tf.range(dim_index), [last_index],\n            tf.range(dim_index + 1, last_index), [dim_index]\n        ], 0), **kwargs)\n\n\n@tf.function\ndef _compute_2d_sparsemax(logits, name=None):\n    \"\"\"Performs the sparsemax operation when axis=-1.\"\"\"\n    shape_op = tf.shape(logits)\n    obs = tf.math.reduce_prod(shape_op[:-1])\n    dims = shape_op[-1]\n\n    # In the paper, they call the logits z.\n    # The mean(logits) can be substracted from logits to make the algorithm\n    # more numerically stable. the instability in this algorithm comes mostly\n    # from the z_cumsum. Substacting the mean will cause z_cumsum to be close\n    # to zero. However, in practise the numerical instability issues are very\n    # minor and substacting the mean causes extra issues with inf and nan\n    # input.\n    # Reshape to [obs, dims] as it is almost free and means the remanining\n    # code doesn't need to worry about the rank.\n    z = tf.reshape(logits, [obs, dims])\n\n    # sort z\n    z_sorted, _ = tf.nn.top_k(z, k=dims)\n\n    # calculate k(z)\n    z_cumsum = tf.math.cumsum(z_sorted, axis=-1)\n    k = tf.range(1, tf.cast(dims, logits.dtype) + 1, dtype=logits.dtype)\n    z_check = 1 + k * z_sorted > z_cumsum\n    # because the z_check vector is always [1,1,...1,0,0,...0] finding the\n    # (index + 1) of the last `1` is the same as just summing the number of 1.\n    k_z = tf.math.reduce_sum(tf.cast(z_check, tf.int32), axis=-1)\n\n    # calculate tau(z)\n    # If there are inf values or all values are -inf, the k_z will be zero,\n    # this is mathematically invalid and will also cause the gather_nd to fail.\n    # Prevent this issue for now by setting k_z = 1 if k_z = 0, this is then\n    # fixed later (see p_safe) by returning p = nan. This results in the same\n    # behavior as softmax.\n    k_z_safe = tf.math.maximum(k_z, 1)\n    indices = tf.stack(\n        [tf.range(0, obs), tf.reshape(k_z_safe, [-1]) - 1], axis=1)\n    tau_sum = tf.gather_nd(z_cumsum, indices)\n    tau_z = (tau_sum - 1) / tf.cast(k_z, logits.dtype)\n\n    # calculate p\n    p = tf.math.maximum(\n        tf.cast(0, logits.dtype), z - tf.expand_dims(tau_z, -1))\n    # If k_z = 0 or if z = nan, then the input is invalid\n    p_safe = tf.where(\n        tf.expand_dims(\n            tf.math.logical_or(\n                tf.math.equal(k_z, 0), tf.math.is_nan(z_cumsum[:, -1])),\n            axis=-1), tf.fill([obs, dims], tf.cast(float(\"nan\"),\n                                                   logits.dtype)), p)\n\n    # Reshape back to original size\n    p_safe = tf.reshape(p_safe, shape_op, name=name)\n    return p_safe\n", "path": "tensorflow_addons/activations/sparsemax.py"}]}