class_docstr: Categorical distribution parametrized by logits.
code: def _unflatten_inputs(self, inputs):
      return jnp.reshape(
          inputs,
          inputs.shape[:-1] + self._shape + (self._n_categories,),
      )
  
imports: import gin
  import gym
  import numpy
  from trax import layers
  from trax.fastmath import numpy
label_desc: Use image node_id 3 for calling the Categorical obj's underlying member method code with example usage: obj._unflatten_inputs(inputs) and returns: jnp