code: def logpmf(
      k: ArrayLike, n: ArrayLike, p: ArrayLike, loc: ArrayLike = 0
  ) -> Array:
      """JAX implementation of scipy.stats.binom.logpmf."""
      k, n, p, loc = promote_args_inexact("binom.logpmf", k, n, p, loc)
      y = lax.sub(k, loc)
      comb_term = lax.sub(
          gammaln(n + 1), lax.add(gammaln(y + 1), gammaln(n - y + 1))
      )
      log_linear_term = lax.add(
          xlogy(y, p), xlog1py(lax.sub(n, y), lax.neg(p))
      )
      log_probs = lax.add(comb_term, log_linear_term)
      return jnp.where(
          lax.ge(k, loc) & lax.lt(k, loc + n + 1), log_probs, -jnp.inf
      )
  
file_path: jax/jax/_src/scipy/stats/binom.py
imports: import scipy.stats
  from jax import lax
  import jax.numpy
  from jax._src.numpy.util import _wraps, promote_args_inexact
  from jax._src.scipy.special import gammaln, xlogy, xlog1py
  from jax._src.typing import Array, ArrayLike
label_desc: Use image node_id 1 for calling a global function with example usage: logpmf(k, n, p, loc) and returns: jnp