code: def _poisson_jax(
      shape, lam, dtype=np.float32, seed=None, name=None, max_iters=None
  ):  # pylint: disable=unused-argument
      """Jax Poisson random sampler."""
      # TODO(b/146674643): use transformed rejection sampling with lam > 10.
      lam = np.array(lam)
      return _poisson_jax_impl(
          lam, seed, _ensure_shape_tuple(shape), dtype, name, max_iters
      )
  
imports: import numpy
  import numpy
  from tensorflow_probability.python.internal.backend.numpy import _utils
  from tensorflow_probability.python.internal.backend.numpy import ops
  from tensorflow_probability.python.internal.backend.numpy.numpy_math import softmax
label_desc: Use image node_id 11 for calling a global function with example usage: _poisson_jax(shape, lam, dtype, seed, name, max_iters) and returns: _poisson_jax_impl