code: def multinomial(
      population_size: int,
      num_samples: int,
      /,
      *,
      batch_size: int = 1,
      probs: Optional[JaxArray] = None,
      replace: bool = True,
      device: jaxlib.xla_extension.Device = None,
      seed: Optional[int] = None,
      out: Optional[JaxArray] = None,
  ) -> JaxArray:
      RNG_, rng_input = jax.random.split(_getRNG())
      _setRNG(RNG_)
      if seed:
          rng_input = jax.random.PRNGKey(seed)
      else:
          RNG_, rng_input = jax.random.split(_getRNG())
          _setRNG(RNG_)
  
      if probs is None:
          probs = (
              jnp.ones(
                  (
                      batch_size,
                      population_size,
                  )
              )
              / population_size
          )
      orig_probs_shape = list(probs.shape)
      num_classes = orig_probs_shape[-1]
      probs_flat = jnp.reshape(probs, (-1, orig_probs_shape[-1]))
      probs_flat = probs_flat / jnp.sum(
          probs_flat, -1, keepdims=True, dtype="float64"
      )
      probs_stack = jnp.split(probs_flat, probs_flat.shape[0])
      samples_stack = [
          jax.random.choice(
              rng_input, num_classes, (num_samples,), replace, p=prob[0]
          )
          for prob in probs_stack
      ]
      samples_flat = jnp.stack(samples_stack)
      return jnp.reshape(
          samples_flat, orig_probs_shape[:-1] + [num_samples]
      )
  
file_path: ivy/ivy/functional/backends/jax/random.py
imports: import jax
  import jax.numpy
  import jaxlib.xla_extension
  from typing import Optional, Union, Sequence
  import ivy
  from ivy.functional.ivy.random import _check_bounds_and_get_shape, _randint_check_dtype_and_bound, _check_valid_scale
  from ivy.functional.backends.jax import JaxArray
  from ivy.func_wrapper import with_unsupported_dtypes
  from .None import backend_version
label_desc: Use image node_id 5 for calling a global function with example usage: multinomial() and returns: jnp