code: def test_temperature_sampler(self):
      tokens = jnp.array([[5, 0, 0, 0]], dtype=jnp.int32)
      cache = None
      key = jax.random.PRNGKey(0)
  
      def tokens_to_logits(tokens, cache):
          jax.debug.print("tokens: {}", tokens)
          logits = jax.nn.one_hot(tokens[..., -1:] + 1, 10)
          logits = jnp.where(logits < 0.5, float("-inf"), logits)
          logits = logits.squeeze(axis=1)
          return logits, cache
  
      new_tokens = temperature_sample(
          tokens, cache, tokens_to_logits, key, topk=5
      )
  
      np.testing.assert_array_equal(new_tokens, [[5, 6, 7, 8]])
  
imports: import jax
  import jax.numpy
  import numpy
  from absl.testing import absltest
  from temperature_sampler import temperature_sample
label_desc: Use image node_id 1 for calling the TestTemperatureSampler obj's underlying member method code with example usage: obj.test_temperature_sampler() and returns: logits,  cache