code: def multinomial(rng, logits, num_samples):
      """Draws samples from a multinomial distribution.
  
      DEPRECATED: Use ``jax.random.categorical`` instead.
  
      Args:
        rng: A JAX PRNGKey.
        logits: Unnormalized log-probabilities, where last dimension is categories.
        num_samples: Number of samples to draw.
  
      Returns:
        Chosen categories, of shape ``logits.shape[:-1] + (num_samples,)``.
      """
      if num_samples != 1:
          shape = (num_samples,) + logits.shape[:-1]
      else:
          shape = None
  
      samples = jax.random.categorical(rng, logits, shape=shape)
  
      # Return expected shape.
      if num_samples != 1:
          return jnp.moveaxis(samples, 0, -1)
      else:
          return samples[..., None]
  
imports: from collections.abc import Iterable, Sequence
  import functools
  from typing import Any, Callable, Optional
  from haiku._src import base
  from haiku._src import initializers
  from haiku._src import module
  from haiku._src import typing
  from haiku._src.typing import PRNGKey
  import jax
  from jax import lax
  import jax.numpy
  import numpy
label_desc: Use image node_id 2 for calling a global function with example usage: multinomial(rng, logits, num_samples) and returns: jnp, samples