prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  Investigate top-p performance
  Currently our top-p (nucleus) sampler has the worst performance of our entire sampling offering, likely because of the sort operation we need to do for our softmax probabilities. See [these rough numbers](https://github.com/keras-team/keras-nlp/commit/41af01200c5d9898f381477b6a67a87b6951f844#diff-576c27794bae284efc278336a16d0649aaf599973a64d4b65dc67bc99bb7f000) for example, where top-p is slower than beam, despite needing much less computation on the model itself.
  
  We should see if there is anything we can do to speed up top-p sampling with XLA (e.g. a more XLA friendly sort op).
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `keras_nlp/samplers/top_p_sampler.py`
  Content:
  ```
  1 # Copyright 2023 The KerasNLP Authors
  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 #     https://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 """Top-p Sampler."""
  15 
  16 import tensorflow as tf
  17 
  18 from keras_nlp.api_export import keras_nlp_export
  19 from keras_nlp.samplers.sampler import Sampler
  20 from keras_nlp.samplers.sampler import call_args_docstring
  21 from keras_nlp.utils.python_utils import format_docstring
  22 
  23 
  24 @format_docstring(call_args=call_args_docstring)
  25 @keras_nlp_export("keras_nlp.samplers.TopPSampler")
  26 class TopPSampler(Sampler):
  27     """Top-P Sampler class.
  28 
  29     This sampler implements top-p search algorithm. Top-p search selects tokens
  30     from the smallest subset of output probabilities that sum to greater than
  31     `p`. Put in another way, top-p will first order token predictions by
  32     likelihood, and ignore all tokens after the cumulative probability of
  33     selected tokens exceeds `p`, then select a token from the remaining tokens.
  34 
  35     Args:
  36         p: float, the `p` value of top-p.
  37         seed: int, defaults to None. The random seed.
  38 
  39     Call Args:
  40         {{call_args}}
  41 
  42     Examples:
  43     ```python
  44     # Use a simple alphabet of lowercase characters with ids in range [0, 25].
  45     int_lookup = {i: chr(i + ord('a')) for i in range(26)}
  46     char_lookup = {v: k for k, v in int_lookup.items()}
  47     batch_size, length, vocab_size = 1, 12, len(int_lookup)
  48 
  49     def next(prompt, cache, index):
  50         hidden_states = tf.ones((batch_size, 10))
  51         # A uniform distribution over our alphabet.
  52         logits = tf.ones((batch_size, vocab_size))
  53         return logits, hidden_states, cache
  54 
  55     output = keras_nlp.samplers.TopPSampler(p=0.1)(
  56         next=next,
  57         prompt=tf.fill((batch_size, length,), char_lookup['z']),
  58         index=5,
  59     )
  60     print(["".join([int_lookup[i] for i in s]) for s in output.numpy()])
  61     # >>> ['zzzzzbabcccb']
  62     ```
  63     """
  64 
  65     def __init__(
  66         self,
  67         p=0.1,
  68         seed=None,
  69     ):
  70         super().__init__()
  71         self.p = p
  72         self.seed = seed
  73 
  74     def get_next_token(self, probabilities):
  75         # Sort preds in descending order.
  76         sorted_preds, sorted_indices = tf.math.top_k(
  77             probabilities, k=tf.shape(probabilities)[1], sorted=True
  78         )
  79         # Calculate cumulative probability distribution.
  80         cumulative_probabilities = tf.math.cumsum(sorted_preds, axis=-1)
  81         # Create a mask for the tokens to keep.
  82         keep_mask = cumulative_probabilities <= self.p
  83         # Shift to include the last token that exceed p.
  84         shifted_keep_mask = tf.concat(
  85             [tf.ones_like(keep_mask[:, :1]), keep_mask[:, :-1]], axis=-1
  86         )
  87         # Filter out unmasked tokens and sample from filtered distribution.
  88         probabilities = tf.where(
  89             shifted_keep_mask,
  90             sorted_preds,
  91             tf.zeros(tf.shape(probabilities), dtype=sorted_preds.dtype),
  92         )
  93         sorted_next_token = tf.random.categorical(
  94             tf.math.log(probabilities), 1, seed=self.seed
  95         )
  96         return tf.gather_nd(sorted_indices, sorted_next_token, batch_dims=1)
  97 
  98     def get_config(self):
  99         config = super().get_config()
  100         config.update(
  101             {
  102                 "p": self.p,
  103                 "seed": self.seed,
  104             }
  105         )
  106         return config
  107
  ```
  --- 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/keras_nlp/samplers/top_p_sampler.py b/keras_nlp/samplers/top_p_sampler.py
  --- a/keras_nlp/samplers/top_p_sampler.py
  +++ b/keras_nlp/samplers/top_p_sampler.py
  @@ -34,6 +34,12 @@
   
       Args:
           p: float, the `p` value of top-p.
  +        k: int, defaults to None. If set, this argument defines a
  +            heuristic "top-k" cutoff applied before the "top-p" sampling. All
  +            logits not in the top `k` will be discarded, and the remaining
  +            logits will be sorted to find a cutoff point for `p`. Setting this
  +            arg can significantly speed sampling up by reducing the number
  +            of tokens to sort.
           seed: int, defaults to None. The random seed.
   
       Call Args:
  @@ -65,16 +71,21 @@
       def __init__(
           self,
           p=0.1,
  +        k=None,
           seed=None,
       ):
           super().__init__()
           self.p = p
  +        self.k = k
           self.seed = seed
   
       def get_next_token(self, probabilities):
  -        # Sort preds in descending order.
  +        cutoff = tf.shape(probabilities)[1]
  +        if self.k is not None:
  +            # If `k` is set, only sample from top `k` tokens.
  +            cutoff = tf.math.minimum(cutoff, self.k)
           sorted_preds, sorted_indices = tf.math.top_k(
  -            probabilities, k=tf.shape(probabilities)[1], sorted=True
  +            probabilities, k=cutoff, sorted=True
           )
           # Calculate cumulative probability distribution.
           cumulative_probabilities = tf.math.cumsum(sorted_preds, axis=-1)
  @@ -88,7 +99,7 @@
           probabilities = tf.where(
               shifted_keep_mask,
               sorted_preds,
  -            tf.zeros(tf.shape(probabilities), dtype=sorted_preds.dtype),
  +            tf.zeros(tf.shape(sorted_preds), dtype=sorted_preds.dtype),
           )
           sorted_next_token = tf.random.categorical(
               tf.math.log(probabilities), 1, seed=self.seed
  @@ -100,6 +111,7 @@
           config.update(
               {
                   "p": self.p,
  +                "k": self.k,
                   "seed": self.seed,
               }
           )
  
verification_info: {"golden_diff": "diff --git a/keras_nlp/samplers/top_p_sampler.py b/keras_nlp/samplers/top_p_sampler.py\n--- a/keras_nlp/samplers/top_p_sampler.py\n+++ b/keras_nlp/samplers/top_p_sampler.py\n@@ -34,6 +34,12 @@\n \n     Args:\n         p: float, the `p` value of top-p.\n+        k: int, defaults to None. If set, this argument defines a\n+            heuristic \"top-k\" cutoff applied before the \"top-p\" sampling. All\n+            logits not in the top `k` will be discarded, and the remaining\n+            logits will be sorted to find a cutoff point for `p`. Setting this\n+            arg can significantly speed sampling up by reducing the number\n+            of tokens to sort.\n         seed: int, defaults to None. The random seed.\n \n     Call Args:\n@@ -65,16 +71,21 @@\n     def __init__(\n         self,\n         p=0.1,\n+        k=None,\n         seed=None,\n     ):\n         super().__init__()\n         self.p = p\n+        self.k = k\n         self.seed = seed\n \n     def get_next_token(self, probabilities):\n-        # Sort preds in descending order.\n+        cutoff = tf.shape(probabilities)[1]\n+        if self.k is not None:\n+            # If `k` is set, only sample from top `k` tokens.\n+            cutoff = tf.math.minimum(cutoff, self.k)\n         sorted_preds, sorted_indices = tf.math.top_k(\n-            probabilities, k=tf.shape(probabilities)[1], sorted=True\n+            probabilities, k=cutoff, sorted=True\n         )\n         # Calculate cumulative probability distribution.\n         cumulative_probabilities = tf.math.cumsum(sorted_preds, axis=-1)\n@@ -88,7 +99,7 @@\n         probabilities = tf.where(\n             shifted_keep_mask,\n             sorted_preds,\n-            tf.zeros(tf.shape(probabilities), dtype=sorted_preds.dtype),\n+            tf.zeros(tf.shape(sorted_preds), dtype=sorted_preds.dtype),\n         )\n         sorted_next_token = tf.random.categorical(\n             tf.math.log(probabilities), 1, seed=self.seed\n@@ -100,6 +111,7 @@\n         config.update(\n             {\n                 \"p\": self.p,\n+                \"k\": self.k,\n                 \"seed\": self.seed,\n             }\n         )\n", "issue": "Investigate top-p performance\nCurrently our top-p (nucleus) sampler has the worst performance of our entire sampling offering, likely because of the sort operation we need to do for our softmax probabilities. See [these rough numbers](https://github.com/keras-team/keras-nlp/commit/41af01200c5d9898f381477b6a67a87b6951f844#diff-576c27794bae284efc278336a16d0649aaf599973a64d4b65dc67bc99bb7f000) for example, where top-p is slower than beam, despite needing much less computation on the model itself.\r\n\r\nWe should see if there is anything we can do to speed up top-p sampling with XLA (e.g. a more XLA friendly sort op).\n", "before_files": [{"content": "# Copyright 2023 The KerasNLP Authors\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#     https://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\"\"\"Top-p Sampler.\"\"\"\n\nimport tensorflow as tf\n\nfrom keras_nlp.api_export import keras_nlp_export\nfrom keras_nlp.samplers.sampler import Sampler\nfrom keras_nlp.samplers.sampler import call_args_docstring\nfrom keras_nlp.utils.python_utils import format_docstring\n\n\n@format_docstring(call_args=call_args_docstring)\n@keras_nlp_export(\"keras_nlp.samplers.TopPSampler\")\nclass TopPSampler(Sampler):\n    \"\"\"Top-P Sampler class.\n\n    This sampler implements top-p search algorithm. Top-p search selects tokens\n    from the smallest subset of output probabilities that sum to greater than\n    `p`. Put in another way, top-p will first order token predictions by\n    likelihood, and ignore all tokens after the cumulative probability of\n    selected tokens exceeds `p`, then select a token from the remaining tokens.\n\n    Args:\n        p: float, the `p` value of top-p.\n        seed: int, defaults to None. The random seed.\n\n    Call Args:\n        {{call_args}}\n\n    Examples:\n    ```python\n    # Use a simple alphabet of lowercase characters with ids in range [0, 25].\n    int_lookup = {i: chr(i + ord('a')) for i in range(26)}\n    char_lookup = {v: k for k, v in int_lookup.items()}\n    batch_size, length, vocab_size = 1, 12, len(int_lookup)\n\n    def next(prompt, cache, index):\n        hidden_states = tf.ones((batch_size, 10))\n        # A uniform distribution over our alphabet.\n        logits = tf.ones((batch_size, vocab_size))\n        return logits, hidden_states, cache\n\n    output = keras_nlp.samplers.TopPSampler(p=0.1)(\n        next=next,\n        prompt=tf.fill((batch_size, length,), char_lookup['z']),\n        index=5,\n    )\n    print([\"\".join([int_lookup[i] for i in s]) for s in output.numpy()])\n    # >>> ['zzzzzbabcccb']\n    ```\n    \"\"\"\n\n    def __init__(\n        self,\n        p=0.1,\n        seed=None,\n    ):\n        super().__init__()\n        self.p = p\n        self.seed = seed\n\n    def get_next_token(self, probabilities):\n        # Sort preds in descending order.\n        sorted_preds, sorted_indices = tf.math.top_k(\n            probabilities, k=tf.shape(probabilities)[1], sorted=True\n        )\n        # Calculate cumulative probability distribution.\n        cumulative_probabilities = tf.math.cumsum(sorted_preds, axis=-1)\n        # Create a mask for the tokens to keep.\n        keep_mask = cumulative_probabilities <= self.p\n        # Shift to include the last token that exceed p.\n        shifted_keep_mask = tf.concat(\n            [tf.ones_like(keep_mask[:, :1]), keep_mask[:, :-1]], axis=-1\n        )\n        # Filter out unmasked tokens and sample from filtered distribution.\n        probabilities = tf.where(\n            shifted_keep_mask,\n            sorted_preds,\n            tf.zeros(tf.shape(probabilities), dtype=sorted_preds.dtype),\n        )\n        sorted_next_token = tf.random.categorical(\n            tf.math.log(probabilities), 1, seed=self.seed\n        )\n        return tf.gather_nd(sorted_indices, sorted_next_token, batch_dims=1)\n\n    def get_config(self):\n        config = super().get_config()\n        config.update(\n            {\n                \"p\": self.p,\n                \"seed\": self.seed,\n            }\n        )\n        return config\n", "path": "keras_nlp/samplers/top_p_sampler.py"}], "after_files": [{"content": "# Copyright 2023 The KerasNLP Authors\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#     https://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\"\"\"Top-p Sampler.\"\"\"\n\nimport tensorflow as tf\n\nfrom keras_nlp.api_export import keras_nlp_export\nfrom keras_nlp.samplers.sampler import Sampler\nfrom keras_nlp.samplers.sampler import call_args_docstring\nfrom keras_nlp.utils.python_utils import format_docstring\n\n\n@format_docstring(call_args=call_args_docstring)\n@keras_nlp_export(\"keras_nlp.samplers.TopPSampler\")\nclass TopPSampler(Sampler):\n    \"\"\"Top-P Sampler class.\n\n    This sampler implements top-p search algorithm. Top-p search selects tokens\n    from the smallest subset of output probabilities that sum to greater than\n    `p`. Put in another way, top-p will first order token predictions by\n    likelihood, and ignore all tokens after the cumulative probability of\n    selected tokens exceeds `p`, then select a token from the remaining tokens.\n\n    Args:\n        p: float, the `p` value of top-p.\n        k: int, defaults to None. If set, this argument defines a\n            heuristic \"top-k\" cutoff applied before the \"top-p\" sampling. All\n            logits not in the top `k` will be discarded, and the remaining\n            logits will be sorted to find a cutoff point for `p`. Setting this\n            arg can significantly speed sampling up by reducing the number\n            of tokens to sort.\n        seed: int, defaults to None. The random seed.\n\n    Call Args:\n        {{call_args}}\n\n    Examples:\n    ```python\n    # Use a simple alphabet of lowercase characters with ids in range [0, 25].\n    int_lookup = {i: chr(i + ord('a')) for i in range(26)}\n    char_lookup = {v: k for k, v in int_lookup.items()}\n    batch_size, length, vocab_size = 1, 12, len(int_lookup)\n\n    def next(prompt, cache, index):\n        hidden_states = tf.ones((batch_size, 10))\n        # A uniform distribution over our alphabet.\n        logits = tf.ones((batch_size, vocab_size))\n        return logits, hidden_states, cache\n\n    output = keras_nlp.samplers.TopPSampler(p=0.1)(\n        next=next,\n        prompt=tf.fill((batch_size, length,), char_lookup['z']),\n        index=5,\n    )\n    print([\"\".join([int_lookup[i] for i in s]) for s in output.numpy()])\n    # >>> ['zzzzzbabcccb']\n    ```\n    \"\"\"\n\n    def __init__(\n        self,\n        p=0.1,\n        k=None,\n        seed=None,\n    ):\n        super().__init__()\n        self.p = p\n        self.k = k\n        self.seed = seed\n\n    def get_next_token(self, probabilities):\n        cutoff = tf.shape(probabilities)[1]\n        if self.k is not None:\n            # If `k` is set, only sample from top `k` tokens.\n            cutoff = tf.math.minimum(cutoff, self.k)\n        sorted_preds, sorted_indices = tf.math.top_k(\n            probabilities, k=cutoff, sorted=True\n        )\n        # Calculate cumulative probability distribution.\n        cumulative_probabilities = tf.math.cumsum(sorted_preds, axis=-1)\n        # Create a mask for the tokens to keep.\n        keep_mask = cumulative_probabilities <= self.p\n        # Shift to include the last token that exceed p.\n        shifted_keep_mask = tf.concat(\n            [tf.ones_like(keep_mask[:, :1]), keep_mask[:, :-1]], axis=-1\n        )\n        # Filter out unmasked tokens and sample from filtered distribution.\n        probabilities = tf.where(\n            shifted_keep_mask,\n            sorted_preds,\n            tf.zeros(tf.shape(sorted_preds), dtype=sorted_preds.dtype),\n        )\n        sorted_next_token = tf.random.categorical(\n            tf.math.log(probabilities), 1, seed=self.seed\n        )\n        return tf.gather_nd(sorted_indices, sorted_next_token, batch_dims=1)\n\n    def get_config(self):\n        config = super().get_config()\n        config.update(\n            {\n                \"p\": self.p,\n                \"k\": self.k,\n                \"seed\": self.seed,\n            }\n        )\n        return config\n", "path": "keras_nlp/samplers/top_p_sampler.py"}]}