class_docstr: Categorical distribution parametrized by logits.
code: def sample(self, inputs, temperature=1.0):
      # No need for LogSoftmax with sampling - softmax normalization is
      # subtracting a constant from every logit, and sampling is taking
      # a max over logits plus noise, so invariant to adding a constant.
      if temperature == 0.0:
          return jnp.argmax(self._unflatten_inputs(inputs), axis=-1)
      return tl.logsoftmax_sample(
          self._unflatten_inputs(inputs), temperature
      )
  
imports: import gin
  import gym
  import numpy
  from trax import layers
  from trax.fastmath import numpy
label_desc: Use image node_id 4 for calling the Categorical obj's underlying member method code with example usage: obj.sample(inputs, temperature) and returns: tl, jnp