code: def logpmf(x: ArrayLike, n: ArrayLike, p: ArrayLike) -> Array:
      """JAX implementation of scipy.stats.multinomial.logpmf."""
      (p,) = promote_args_inexact("multinomial.logpmf", p)
      x, n = promote_args_numeric("multinomial.logpmf", x, n)
      if not jnp.issubdtype(x.dtype, jnp.integer):
          raise ValueError(
              f"x and n must be of integer type; got x.dtype={x.dtype}, n.dtype={n.dtype}"
          )
      x = x.astype(p.dtype)
      n = n.astype(p.dtype)
      logprobs = gammaln(n + 1) + jnp.sum(
          xlogy(x, p) - gammaln(x + 1), axis=-1
      )
      return jnp.where(jnp.equal(jnp.sum(x), n), logprobs, -jnp.inf)
  
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 1 for calling a global function with example usage: logpmf(x, n, p) and returns: jnp