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