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.


Discussion:
`scipy.sparse.lil_matrix` is not a JAX-compatible type, so it cannot be used as a return value in jit-compiled code. Your code should work fine if you remove `@partial(jit, static_argnums=(0, 1,))`; if you want to use JAX jit, you'll have to rewrite the jitted parts of your code to avoid the use of `scipy.sparse` objects.
That said, that's a pretty bad error message, and I'm not entirely sure how to trigger it. Can you share a small reproduction of the `vector::reserve` error?
I managed to reproduce it from the above code this way:
```python
import pandas as pd
from scipy.sparse import lil_matrix, csr_matrix
import numpy as np
from jax import jit, partial
from typing import Dict

def optimize(x):
  return x

@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

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))

_get_pairwise_frequencies(
    pd.DataFrame({'x': [1, 2, 3], 'y': [4, 5, 6]})
)
```
I'm certain it could be simplified :grin:
It seems the key to the repro is to return an invalid jax type while passing a static argument by keyword:
```python
from scipy.sparse import lil_matrix
from jax import jit, partial

@partial(jit, static_argnums=0)
def f(size) -> lil_matrix:
  return lil_matrix((size, size))

f(size=4)
```
```pytb
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-57-eb6a7f1d5936> in <module>()
      6   return lil_matrix((size, size))
      7 
----> 8 f(size=4)

/usr/local/lib/python3.6/dist-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
```
If you instead pass the static argument by position, you get a more reasonable error:
```python
f(4)
```
```pytb
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-58-adf85db3d866> in <module>()
      6   return lil_matrix((size, size))
      7 
----> 8 f(4)

-----[15 frames]

/usr/local/lib/python3.6/dist-packages/jax/core.py in concrete_aval(x)
    860     handler = pytype_aval_mappings.get(typ)
    861     if handler: return handler(x)
--> 862   raise TypeError(f"{type(x)} is not a valid JAX type")
    863 
    864 

TypeError: <class 'scipy.sparse.lil.lil_matrix'> is not a valid JAX type
```
Actually, it turns out `vector::reserve` comes from any function passing a static arg by keyword:
```python
from jax import jit
f = jit(lambda x: x, static_argnums=0)
f(x=4)
```
```pytb
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-65-ed41108eed0c> in <module>()
      1 from jax import jit
      2 f = jit(lambda x: x, static_argnums=0)
----> 3 f(x=4)

/usr/local/lib/python3.6/dist-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
```
We should either handle the keyword argument appropriately in jitted functions, or if that is difficult, provide a better error to the user.

Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -266,9 +266,9 @@ def cache_miss(_, *args, **kwargs):
     # (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)


Test patch: