Repository: google/jax
Issue: How to resolve ValueError `vector::reserve`?

I am trying to optimize the following function using jit:

```py
@partial(jit, static_argnums=(0, 1,))
def coocurrence_helper(pairs: np.array, label_map: Dict) -> lil_matrix:
    uniques = lil_matrix(np.zeros((len(label_map), len(label_map))).astype("int32"))
    for item in pairs:
        if item[0]!=item[1]:
            uniques[label_map[item[0]], label_map[item[1]]] += 1
    return uniques
```

the routine above is used here:

```py
def _get_pairwise_frequencies(
     data: pd.DataFrame, crosstab=False
    ) -> pd.DataFrame:
        values = data.stack()
        values.index = values.index.droplevel(1)
        values.name = "vals"
        values = optimize(values.to_frame())
        pair = optimize(values.join(values, rsuffix="_2"))
        label_map = dict()
        for lbl, each in enumerate(values.vals.unique()):
            label_map[each] = lbl
        if not crosstab:
            freq = coocurrence_helper(pairs = pair.values, label_map=label_map)
            return ((freq / freq.sum(1).ravel()).astype(np.float32))
        else:
            freq = pd.crosstab(pair["vals"], pair["vals_2"])
            self.index = freq.index
            return csr_matrix((freq / freq.sum(1)).astype(np.float32))
```

But i get the following error:

```py
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-42-f8e638fc2bb6> in <module>
----> 1 _get_pairwise_frequencies(data)

<ipython-input-30-43adeb39c76c> in _get_pairwise_frequencies(data, crosstab)
     25             label_map[each] = lbl
     26         if not crosstab:
---> 27             freq = coocurrence_helper(pairs = pair.values, label_map=label_map)
     28             return csr_matrix((freq / freq.sum(1).ravel()).astype(np.float32))
     29         else:

~/anaconda3/envs/tensorflow2_p36/lib/python3.6/site-packages/jax/api.py in f_jitted(*args, **kwargs)
    369         return cache_miss(*args, **kwargs)[0]  # probably won't return
    370     else:
--> 371       return cpp_jitted_f(*args, **kwargs)
    372   f_jitted._cpp_jitted_f = cpp_jitted_f
    373 

ValueError: vector::reserve
```

What can be the source of the issue here? Without using `static_argnums` the error message is 
```py
RuntimeError: Invalid argument: Unknown NumPy type O size 8
```
with the same traceback.

Changed files:
jax/api.py

diff --git a/jax/api.py b/jax/api.py
index 9f450c3ef..f569b9e61 100644
--- a/jax/api.py
+++ b/jax/api.py
@@ -266,9 +266,9 @@ def _cpp_jit(
     # (dyn_args, donated_invars, args_flat, in_tree), since otherwise we have
     # work/code that is redundant between C++ and Python. We can try that later.
     if max(static_argnums + donate_argnums, default=-1) >= len(args):
-      raise ValueError(f"jitted function has static_argnums={static_argnums}, "
-                       f"donate_argnums={donate_argnums} but "
-                       f"was called with only {len(args)} positional arguments.")
+      msg = ("jitted function has static_argnums={}, donate_argnums={} but "
+             "was called with only {} positional arguments.")
+      raise ValueError(msg.format(static_argnums, donate_argnums, len(args)))
     f = lu.wrap_init(fun)
     if static_argnums:
       f, dyn_args = argnums_partial_except(f, static_argnums, args)