How can I generate an array of samples from a poisson distribution in python using jax (jit)?

I am using jax (https://github.com/google/jax) to code up a neural network, and to simulate my inputs I would like to generate an array of samples from a poisson distribution. How can I do this given jax's restrictions?

I have already tried using np.random.poisson(mu,N) and scipy.stats.poisson.rvs(mu, size=N). Neither of these work because they are not supported in jax.numpy and jax.scipy.stats. So basically, I need an alternative solution, either with another package that jax does support or by hard-coding the poisson function.

from jax import jit, vmap
import jax.numpy as np
import numpy as onp
from scipy.stats import poisson

def build_input_and_targets_simulated(ntime, key):
    """
    Function: Simulate inputs and targets.
    Args:
        ntime: number of time steps in input
        key: key for random number generator
    Returns:
        inputs: txu matrix of inputs
        targets: txu matrix of target classifications
    """
    mu = 0.03  # average number of events per interval

    # scipy method
    inputs = np.array([poisson.rvs(mu, size=ntime), poisson.rvs(mu, 
size=ntime)]).T

    # numpy method
    inputs = np.array([onp.random.poisson(mu, ntime), 
onp.random.poisson(mu, ntime)]).T

    targets = onp.zeros((ntime,1))

    # determine target based on difference in inputs
    diffT = np.cumsum(inputs[:,0]) - np.cumsum(inputs[:,1]) # calculate 
cumulative difference in inputs for each time point

    targets[diffT > 0] = 1 # if diffT > 0, set binary choice
    targets[diffT ==0] = int(random.randint(key, (1,1), 0, 2)) # if 
inputs are equal, select target class randomly

    return inputs, targets

# Now batch it and jit.
build_input_and_target = build_input_and_targets_simulated
build_inputs_and_targets = vmap(build_input_and_target, in_axes=(None, 
0))
build_inputs_and_targets_jit = jit(build_inputs_and_targets, 
static_argnums=(0,))


seed = onp.random.randint(0, 1000000)
key = random.PRNGKey(seed)
ntimesteps = 25

inputs, targets = build_inputs_and_targets_jit(ntimesteps, key)

If I use the scipy method, I get an error that looks like this:

Exception: Tracer can't be used with raw numpy functions. You might have
  import numpy as np
instead of
  import jax.numpy as np

If I use the numpy method, I get an error that looks like this:

TypeError: 'BatchTracer' object cannot be interpreted as an integer

Both of these errors appear to be related to the specialized data types that jax uses (and which are necessary when using jit).

How can I get around this??