class_docstr: Categorical distribution parametrized by logits.
code: def n_inputs(self):
      return np.prod(self._shape, dtype=jnp.int32) * 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 2 for calling the Categorical obj's underlying member method code with example usage: obj.n_inputs() and returns: unknown