code: def pmf(x: ArrayLike, n: ArrayLike, p: ArrayLike) -> Array:
      """JAX implementation of scipy.stats.multinomial.pmf."""
      return lax.exp(logpmf(x, n, p))
  
file_path: jax/jax/_src/scipy/stats/multinomial.py
imports: import scipy.stats
  from jax import lax
  import jax.numpy
  from jax._src.numpy.util import _wraps, promote_args_inexact, promote_args_numeric
  from jax._src.scipy.special import gammaln, xlogy
  from jax._src.typing import Array, ArrayLike
label_desc: Use image node_id 2 for calling a global function with example usage: pmf(x, n, p) and returns: lax