code: def logpmf(k: ArrayLike, p: ArrayLike, loc: ArrayLike = 0) -> Array:
      k, p, loc = promote_args_inexact("geom.logpmf", k, p, loc)
      zero = _lax_const(k, 0)
      one = _lax_const(k, 1)
      x = lax.sub(k, loc)
      log_probs = xlog1py(lax.sub(x, one), -p) + lax.log(p)
      return jnp.where(lax.le(x, zero), -jnp.inf, log_probs)
  
file_path: jax/jax/_src/scipy/stats/geom.py
imports: import scipy.stats
  from jax import lax
  import jax.numpy
  from jax._src.lax.lax import _const
  from jax._src.numpy.util import _wraps, promote_args_inexact
  from jax.scipy.special import 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, p, loc) and returns: jnp