prompt: 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()
  
  ```
  
golden_diff: 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)
  
verification_info: {"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"}]}