We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
lm1b sampler update is out of bound?
Hi,

_Disclaimer: I'm not confident if this is some intended behavior._

I think in the lm1b example, the prediction sampler could write to indices out of bound of the result.
The particular offending line is:  https://github.com/google/flax/blob/master/examples/lm1b/temperature_sampler.py#L109.

Since `i <= max_decode_len` according to https://github.com/google/flax/blob/master/examples/lm1b/temperature_sampler.py#L75, and since `prompt_inputs.shape[1] == max_decode_len`, the last iteration (or two iterations?) of while loop would overwrite to `prompt_inputs[..., max_decode_len - 1]` (due to XLA semantics). 

It's either a subtle bug (it won't raise cause error), or there's something I'm not understanding about the sampling procedure.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `examples/lm1b/temperature_sampler.py`
Content:
```
1 # Copyright 2022 The Flax 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 #     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 """Fast decoding routines for inference from a trained language model."""
16 
17 from jax import lax
18 from jax import random
19 import jax.numpy as jnp
20 
21 
22 # Constants
23 # The default End-of-Sentence token id is 2 (SentencePiece).
24 EOS_ID = 2
25 
26 
27 def temperature_sample(prompt_inputs,
28                        init_cache,
29                        tokens_to_logits,
30                        prng_key,
31                        temperature=1.0,
32                        topk=20,
33                        eos_token=EOS_ID):
34   """Temperature sampling for language model generation.
35 
36   Args:
37     prompt_inputs: array: [batch_size, max_decode_len] int32 sequence of tokens.
38     init_cache: flax attention cache.
39     tokens_to_logits: fast autoregressive decoder function taking single token
40       slices and cache and returning next-token logits and updated cache.
41     prng_key: JAX PRNGKey.
42     temperature: float: sampling temperature factor. As it approaches
43       zero this becomes equivalent to greedy sampling.
44     topk: integer: if nonzero only use the top-k logits to sample next token,
45       if zero don't use any cutoff and sample from full logits over vocabulary.
46     eos_token: int: end-of-sentence token for target vocabulary.
47 
48   Returns:
49     Array of sampled sequences: [batch_size, max_decode_len]
50   """
51   batch_size = prompt_inputs.shape[0]
52   max_decode_len = prompt_inputs.shape[1]
53   end_marker = jnp.array(eos_token)
54   temperature = jnp.array(temperature)
55 
56   # Initialize sampling loop state.
57   # initial loop PRNGKey
58   rng0 = prng_key
59   # loop position counter.
60   i0 = jnp.array(0)
61   # per batch-item holding current token in loop.
62   token0 = jnp.zeros((batch_size, 1), dtype=jnp.int32)
63   # per batch-item state bit indicating if sentence has finished.
64   ended0 = jnp.zeros((batch_size, 1), dtype=jnp.bool_)
65   # (batch, length) array containing prefix prompt tokens for sampling loop
66   # as well as the generated output of newly sampled tokens.
67   sequences0 = prompt_inputs
68   # Sampling loop state is stored in a simple tuple.
69   sampling_loop_init_state = (i0, sequences0, init_cache, token0, ended0, rng0)
70 
71   def sampling_loop_cond_fn(state):
72     """Sampling loop termination condition."""
73     (i, _, _, _, ended, _) = state
74     # Have we reached max decoding length?
75     not_at_end = (i < max_decode_len)
76     # Have all sampled sequences reached an end marker?
77     all_sequences_ended = jnp.all(ended)
78     return not_at_end & (~all_sequences_ended)
79 
80   def sampling_loop_body_fn(state):
81     """Sampling loop state update."""
82     i, sequences, cache, cur_token, ended, rng = state
83     # Split RNG for sampling.
84     rng1, rng2 = random.split(rng)
85     # Call fast-decoder model on current tokens to get next-position logits.
86     logits, new_cache = tokens_to_logits(cur_token, cache)
87     # Sample next token from logits.
88     # TODO(levskaya): add top-p "nucleus" sampling option.
89     if topk:
90       # Get top-k logits and their indices, sample within these top-k tokens.
91       topk_logits, topk_idxs = lax.top_k(logits, topk)
92       topk_token = jnp.expand_dims(random.categorical(
93           rng1, topk_logits / temperature).astype(jnp.int32), axis=-1)
94       # Return the original indices corresponding to the sampled top-k tokens.
95       next_token = jnp.squeeze(
96           jnp.take_along_axis(topk_idxs, topk_token, axis=-1), axis=-1)
97     else:
98       next_token = random.categorical(
99           rng1, logits / temperature).astype(jnp.int32)
100     # Only use sampled tokens if we're past provided prefix tokens.
101     out_of_prompt = (sequences[:, i+1] == 0)
102     next_token = (next_token * out_of_prompt +
103                   sequences[:, i+1] * ~out_of_prompt)
104     # If end-marker reached for batch item, only emit padding tokens.
105     next_token_or_endpad = (next_token[None] * ~ended)
106     ended |= (next_token_or_endpad == end_marker)
107     # Add current sampled tokens to recorded sequences.
108     new_sequences = lax.dynamic_update_slice(
109         sequences, next_token_or_endpad, (0, i+1))
110     return (i+1, new_sequences, new_cache, next_token_or_endpad, ended, rng2)
111 
112   # Run sampling loop and collect final state.
113   final_state = lax.while_loop(sampling_loop_cond_fn,
114                                sampling_loop_body_fn,
115                                sampling_loop_init_state)
116 
117   # Pick part of the state corresponding to the sampled sequences.
118   final_sequences = final_state[1]
119   return final_sequences
120
```
--- 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/examples/lm1b/temperature_sampler.py b/examples/lm1b/temperature_sampler.py
--- a/examples/lm1b/temperature_sampler.py
+++ b/examples/lm1b/temperature_sampler.py
@@ -57,7 +57,7 @@
   # initial loop PRNGKey
   rng0 = prng_key
   # loop position counter.
-  i0 = jnp.array(0)
+  i0 = jnp.array(-1)
   # per batch-item holding current token in loop.
   token0 = jnp.zeros((batch_size, 1), dtype=jnp.int32)
   # per batch-item state bit indicating if sentence has finished.
@@ -72,7 +72,7 @@
     """Sampling loop termination condition."""
     (i, _, _, _, ended, _) = state
     # Have we reached max decoding length?
-    not_at_end = (i < max_decode_len)
+    not_at_end = (i < max_decode_len - 1)
     # Have all sampled sequences reached an end marker?
     all_sequences_ended = jnp.all(ended)
     return not_at_end & (~all_sequences_ended)

{"golden_diff": "diff --git a/examples/lm1b/temperature_sampler.py b/examples/lm1b/temperature_sampler.py\n--- a/examples/lm1b/temperature_sampler.py\n+++ b/examples/lm1b/temperature_sampler.py\n@@ -57,7 +57,7 @@\n   # initial loop PRNGKey\n   rng0 = prng_key\n   # loop position counter.\n-  i0 = jnp.array(0)\n+  i0 = jnp.array(-1)\n   # per batch-item holding current token in loop.\n   token0 = jnp.zeros((batch_size, 1), dtype=jnp.int32)\n   # per batch-item state bit indicating if sentence has finished.\n@@ -72,7 +72,7 @@\n     \"\"\"Sampling loop termination condition.\"\"\"\n     (i, _, _, _, ended, _) = state\n     # Have we reached max decoding length?\n-    not_at_end = (i < max_decode_len)\n+    not_at_end = (i < max_decode_len - 1)\n     # Have all sampled sequences reached an end marker?\n     all_sequences_ended = jnp.all(ended)\n     return not_at_end & (~all_sequences_ended)\n", "issue": "lm1b sampler update is out of bound?\nHi,\r\n\r\n_Disclaimer: I'm not confident if this is some intended behavior._\r\n\r\nI think in the lm1b example, the prediction sampler could write to indices out of bound of the result.\r\nThe particular offending line is:  https://github.com/google/flax/blob/master/examples/lm1b/temperature_sampler.py#L109.\r\n\r\nSince `i <= max_decode_len` according to https://github.com/google/flax/blob/master/examples/lm1b/temperature_sampler.py#L75, and since `prompt_inputs.shape[1] == max_decode_len`, the last iteration (or two iterations?) of while loop would overwrite to `prompt_inputs[..., max_decode_len - 1]` (due to XLA semantics). \r\n\r\nIt's either a subtle bug (it won't raise cause error), or there's something I'm not understanding about the sampling procedure.\n", "before_files": [{"content": "# Copyright 2022 The Flax 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#     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\"\"\"Fast decoding routines for inference from a trained language model.\"\"\"\n\nfrom jax import lax\nfrom jax import random\nimport jax.numpy as jnp\n\n\n# Constants\n# The default End-of-Sentence token id is 2 (SentencePiece).\nEOS_ID = 2\n\n\ndef temperature_sample(prompt_inputs,\n                       init_cache,\n                       tokens_to_logits,\n                       prng_key,\n                       temperature=1.0,\n                       topk=20,\n                       eos_token=EOS_ID):\n  \"\"\"Temperature sampling for language model generation.\n\n  Args:\n    prompt_inputs: array: [batch_size, max_decode_len] int32 sequence of tokens.\n    init_cache: flax attention cache.\n    tokens_to_logits: fast autoregressive decoder function taking single token\n      slices and cache and returning next-token logits and updated cache.\n    prng_key: JAX PRNGKey.\n    temperature: float: sampling temperature factor. As it approaches\n      zero this becomes equivalent to greedy sampling.\n    topk: integer: if nonzero only use the top-k logits to sample next token,\n      if zero don't use any cutoff and sample from full logits over vocabulary.\n    eos_token: int: end-of-sentence token for target vocabulary.\n\n  Returns:\n    Array of sampled sequences: [batch_size, max_decode_len]\n  \"\"\"\n  batch_size = prompt_inputs.shape[0]\n  max_decode_len = prompt_inputs.shape[1]\n  end_marker = jnp.array(eos_token)\n  temperature = jnp.array(temperature)\n\n  # Initialize sampling loop state.\n  # initial loop PRNGKey\n  rng0 = prng_key\n  # loop position counter.\n  i0 = jnp.array(0)\n  # per batch-item holding current token in loop.\n  token0 = jnp.zeros((batch_size, 1), dtype=jnp.int32)\n  # per batch-item state bit indicating if sentence has finished.\n  ended0 = jnp.zeros((batch_size, 1), dtype=jnp.bool_)\n  # (batch, length) array containing prefix prompt tokens for sampling loop\n  # as well as the generated output of newly sampled tokens.\n  sequences0 = prompt_inputs\n  # Sampling loop state is stored in a simple tuple.\n  sampling_loop_init_state = (i0, sequences0, init_cache, token0, ended0, rng0)\n\n  def sampling_loop_cond_fn(state):\n    \"\"\"Sampling loop termination condition.\"\"\"\n    (i, _, _, _, ended, _) = state\n    # Have we reached max decoding length?\n    not_at_end = (i < max_decode_len)\n    # Have all sampled sequences reached an end marker?\n    all_sequences_ended = jnp.all(ended)\n    return not_at_end & (~all_sequences_ended)\n\n  def sampling_loop_body_fn(state):\n    \"\"\"Sampling loop state update.\"\"\"\n    i, sequences, cache, cur_token, ended, rng = state\n    # Split RNG for sampling.\n    rng1, rng2 = random.split(rng)\n    # Call fast-decoder model on current tokens to get next-position logits.\n    logits, new_cache = tokens_to_logits(cur_token, cache)\n    # Sample next token from logits.\n    # TODO(levskaya): add top-p \"nucleus\" sampling option.\n    if topk:\n      # Get top-k logits and their indices, sample within these top-k tokens.\n      topk_logits, topk_idxs = lax.top_k(logits, topk)\n      topk_token = jnp.expand_dims(random.categorical(\n          rng1, topk_logits / temperature).astype(jnp.int32), axis=-1)\n      # Return the original indices corresponding to the sampled top-k tokens.\n      next_token = jnp.squeeze(\n          jnp.take_along_axis(topk_idxs, topk_token, axis=-1), axis=-1)\n    else:\n      next_token = random.categorical(\n          rng1, logits / temperature).astype(jnp.int32)\n    # Only use sampled tokens if we're past provided prefix tokens.\n    out_of_prompt = (sequences[:, i+1] == 0)\n    next_token = (next_token * out_of_prompt +\n                  sequences[:, i+1] * ~out_of_prompt)\n    # If end-marker reached for batch item, only emit padding tokens.\n    next_token_or_endpad = (next_token[None] * ~ended)\n    ended |= (next_token_or_endpad == end_marker)\n    # Add current sampled tokens to recorded sequences.\n    new_sequences = lax.dynamic_update_slice(\n        sequences, next_token_or_endpad, (0, i+1))\n    return (i+1, new_sequences, new_cache, next_token_or_endpad, ended, rng2)\n\n  # Run sampling loop and collect final state.\n  final_state = lax.while_loop(sampling_loop_cond_fn,\n                               sampling_loop_body_fn,\n                               sampling_loop_init_state)\n\n  # Pick part of the state corresponding to the sampled sequences.\n  final_sequences = final_state[1]\n  return final_sequences\n", "path": "examples/lm1b/temperature_sampler.py"}], "after_files": [{"content": "# Copyright 2022 The Flax 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#     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\"\"\"Fast decoding routines for inference from a trained language model.\"\"\"\n\nfrom jax import lax\nfrom jax import random\nimport jax.numpy as jnp\n\n\n# Constants\n# The default End-of-Sentence token id is 2 (SentencePiece).\nEOS_ID = 2\n\n\ndef temperature_sample(prompt_inputs,\n                       init_cache,\n                       tokens_to_logits,\n                       prng_key,\n                       temperature=1.0,\n                       topk=20,\n                       eos_token=EOS_ID):\n  \"\"\"Temperature sampling for language model generation.\n\n  Args:\n    prompt_inputs: array: [batch_size, max_decode_len] int32 sequence of tokens.\n    init_cache: flax attention cache.\n    tokens_to_logits: fast autoregressive decoder function taking single token\n      slices and cache and returning next-token logits and updated cache.\n    prng_key: JAX PRNGKey.\n    temperature: float: sampling temperature factor. As it approaches\n      zero this becomes equivalent to greedy sampling.\n    topk: integer: if nonzero only use the top-k logits to sample next token,\n      if zero don't use any cutoff and sample from full logits over vocabulary.\n    eos_token: int: end-of-sentence token for target vocabulary.\n\n  Returns:\n    Array of sampled sequences: [batch_size, max_decode_len]\n  \"\"\"\n  batch_size = prompt_inputs.shape[0]\n  max_decode_len = prompt_inputs.shape[1]\n  end_marker = jnp.array(eos_token)\n  temperature = jnp.array(temperature)\n\n  # Initialize sampling loop state.\n  # initial loop PRNGKey\n  rng0 = prng_key\n  # loop position counter.\n  i0 = jnp.array(-1)\n  # per batch-item holding current token in loop.\n  token0 = jnp.zeros((batch_size, 1), dtype=jnp.int32)\n  # per batch-item state bit indicating if sentence has finished.\n  ended0 = jnp.zeros((batch_size, 1), dtype=jnp.bool_)\n  # (batch, length) array containing prefix prompt tokens for sampling loop\n  # as well as the generated output of newly sampled tokens.\n  sequences0 = prompt_inputs\n  # Sampling loop state is stored in a simple tuple.\n  sampling_loop_init_state = (i0, sequences0, init_cache, token0, ended0, rng0)\n\n  def sampling_loop_cond_fn(state):\n    \"\"\"Sampling loop termination condition.\"\"\"\n    (i, _, _, _, ended, _) = state\n    # Have we reached max decoding length?\n    not_at_end = (i < max_decode_len - 1)\n    # Have all sampled sequences reached an end marker?\n    all_sequences_ended = jnp.all(ended)\n    return not_at_end & (~all_sequences_ended)\n\n  def sampling_loop_body_fn(state):\n    \"\"\"Sampling loop state update.\"\"\"\n    i, sequences, cache, cur_token, ended, rng = state\n    # Split RNG for sampling.\n    rng1, rng2 = random.split(rng)\n    # Call fast-decoder model on current tokens to get next-position logits.\n    logits, new_cache = tokens_to_logits(cur_token, cache)\n    # Sample next token from logits.\n    # TODO(levskaya): add top-p \"nucleus\" sampling option.\n    if topk:\n      # Get top-k logits and their indices, sample within these top-k tokens.\n      topk_logits, topk_idxs = lax.top_k(logits, topk)\n      topk_token = jnp.expand_dims(random.categorical(\n          rng1, topk_logits / temperature).astype(jnp.int32), axis=-1)\n      # Return the original indices corresponding to the sampled top-k tokens.\n      next_token = jnp.squeeze(\n          jnp.take_along_axis(topk_idxs, topk_token, axis=-1), axis=-1)\n    else:\n      next_token = random.categorical(\n          rng1, logits / temperature).astype(jnp.int32)\n    # Only use sampled tokens if we're past provided prefix tokens.\n    out_of_prompt = (sequences[:, i+1] == 0)\n    next_token = (next_token * out_of_prompt +\n                  sequences[:, i+1] * ~out_of_prompt)\n    # If end-marker reached for batch item, only emit padding tokens.\n    next_token_or_endpad = (next_token[None] * ~ended)\n    ended |= (next_token_or_endpad == end_marker)\n    # Add current sampled tokens to recorded sequences.\n    new_sequences = lax.dynamic_update_slice(\n        sequences, next_token_or_endpad, (0, i+1))\n    return (i+1, new_sequences, new_cache, next_token_or_endpad, ended, rng2)\n\n  # Run sampling loop and collect final state.\n  final_state = lax.while_loop(sampling_loop_cond_fn,\n                               sampling_loop_body_fn,\n                               sampling_loop_init_state)\n\n  # Pick part of the state corresponding to the sampled sequences.\n  final_sequences = final_state[1]\n  return final_sequences\n", "path": "examples/lm1b/temperature_sampler.py"}]}