# Buggy code
def choice(key, a, shape=(), replace=True, p=None):
    """Generates a random sample from a given 1-D array.

    Args:
      key: a PRNGKey used as the random key.
      a : 1D array or int. If an ndarray, a random sample is generated from
        its elements. If an int, the random sample is generated as if a were
        arange(a).
      shape : tuple of ints, optional. Output shape.  If the given shape is,
        e.g., ``(m, n)``, then ``m * n`` samples are drawn.  Default is (),
        in which case a single value is returned.
      replace : boolean.  Whether the sample is with or without replacement.
        default is True.
      p : 1-D array-like, The probabilities associated with each entry in a.
        If not given the sample assumes a uniform distribution over all
        entries in a.

    Returns:
      An array of shape `shape` containing samples from `a`.
    """
    a = jnp.asarray(a)
    if a.ndim not in [0, 1]:
        raise ValueError("a must be an integer or 1-dimensional")
    n_inputs = int(a) if a.ndim == 0 else len(a)
    n_draws = prod(shape)
    if n_draws == 0:
        return jnp.zeros(shape, dtype=a.dtype)
    if n_inputs <= 0:
        raise ValueError("a must be greater than 0 unless no samples are taken")
    if not replace and n_draws > n_inputs:
        raise ValueError(
            "Cannot take a larger sample than population when 'replace=False'"
        )

    if p is None:
        if replace:
            ind = randint(key, shape, 0, n_inputs)
            result = ind if a.ndim == 0 else a[ind]
        else:
            result = permutation(key, a)[:n_draws]
    else:
        p = jnp.asarray(p)
        if p.shape != (n_inputs,):
            raise ValueError("p must be None or match the shape of a")
        if replace:
            p_cuml = jnp.cumsum(p)
            r = p_cuml[-1] * (1 - uniform(key, shape))
            ind = jnp.searchsorted(p_cuml, r)
            result = ind if a.ndim == 0 else a[ind]
        else:
            # Gumbel top-k trick: https://timvieira.github.io/blog/post/2019/09/16/algorithms-for-sampling-without-replacement/
            g = -gumbel(key, (n_inputs,)) - jnp.log(p)
            ind = jnp.argsort(g)[:n_draws]
            result = ind if a.ndim == 0 else a[ind]
    return result.reshape(shape)


# Traceback (TypeError)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-19-56696338bc13> in <module>()
1 key = jax.random.PRNGKey(0)
2 jax.random.choice(key, 5, 2, replace=False)
----> 3 jax.random.choice(key, 5, 2, replace=True)

2 frames
/usr/local/lib/python3.6/dist-packages/jax/random.py in choice(key, a, shape, replace, p)
566   if p is None:
567     if replace:
--> 568       ind = randint(key, shape, 0, n_inputs)
569       result = ind if a.ndim == 0 else a[ind]
570     else:

/usr/local/lib/python3.6/dist-packages/jax/random.py in randint(key, shape, minval, maxval, dtype)
417   """
418   dtype = dtypes.canonicalize_dtype(dtype)
--> 419   shape = abstract_arrays.canonicalize_shape(shape)
420   return _randint(key, shape, minval, maxval, dtype)
421

/usr/local/lib/python3.6/dist-packages/jax/core.py in canonicalize_shape(shape)
1078          "got {}.")
1079   if any(isinstance(x, Tracer) and isinstance(get_aval(x), ShapedArray)
-> 1080          and not isinstance(get_aval(x), ConcreteArray) for x in shape):
1081     msg += ("\nIf using `jit`, try using `static_argnums` or applying `jit` to "
1082             "smaller subfunctions.")

TypeError: 'int' object is not iterable

# Fixed code
def choice(key, a, shape=(), replace=True, p=None):
    """Generates a random sample from a given 1-D array.

    Args:
      key: a PRNGKey used as the random key.
      a : 1D array or int. If an ndarray, a random sample is generated from
        its elements. If an int, the random sample is generated as if a were
        arange(a).
      shape : tuple of ints, optional. Output shape.  If the given shape is,
        e.g., ``(m, n)``, then ``m * n`` samples are drawn.  Default is (),
        in which case a single value is returned.
      replace : boolean.  Whether the sample is with or without replacement.
        default is True.
      p : 1-D array-like, The probabilities associated with each entry in a.
        If not given the sample assumes a uniform distribution over all
        entries in a.

    Returns:
      An array of shape `shape` containing samples from `a`.
    """
    if not isinstance(shape, Sequence):
        raise TypeError(
            f"shape argument of jax.random.choice must be a sequence, got {shape}"
        )
    a = jnp.asarray(a)
    if a.ndim not in [0, 1]:
        raise ValueError("a must be an integer or 1-dimensional")
    n_inputs = int(a) if a.ndim == 0 else len(a)
    n_draws = prod(shape)
    if n_draws == 0:
        return jnp.zeros(shape, dtype=a.dtype)
    if n_inputs <= 0:
        raise ValueError("a must be greater than 0 unless no samples are taken")
    if not replace and n_draws > n_inputs:
        raise ValueError(
            "Cannot take a larger sample than population when 'replace=False'"
        )

    if p is None:
        if replace:
            ind = randint(key, shape, 0, n_inputs)
            result = ind if a.ndim == 0 else a[ind]
        else:
            result = permutation(key, a)[:n_draws]
    else:
        p = jnp.asarray(p)
        if p.shape != (n_inputs,):
            raise ValueError("p must be None or match the shape of a")
        if replace:
            p_cuml = jnp.cumsum(p)
            r = p_cuml[-1] * (1 - uniform(key, shape))
            ind = jnp.searchsorted(p_cuml, r)
            result = ind if a.ndim == 0 else a[ind]
        else:
            # Gumbel top-k trick: https://timvieira.github.io/blog/post/2019/09/16/algorithms-for-sampling-without-replacement/
            g = -gumbel(key, (n_inputs,)) - jnp.log(p)
            ind = jnp.argsort(g)[:n_draws]
            result = ind if a.ndim == 0 else a[ind]
    return result.reshape(shape)