code: def softmax_sample(key, logits, _, temp=1):
      return (
          jax.random.categorical(key, logits / temp, -1).astype(
              jnp.uint32
          ),
          None,
      )
  
imports: import jax
  import jax.numpy
label_desc: Use image node_id 1 for calling a global function with example usage: softmax_sample(key, logits, _, temp) and returns: jax,  None