Issue 1:
Currently, if a static_argnums or static_argname doesn't exist for the function, nothing happens

This can be an issue if a function is updated but the jax.jit not being updated

import jax
import functools

@functools.partial(jax.jit, static_argnames=('b', 'd'))
def f(a, b, c):
return a + b + c

print(f(1., 2., 3.))

@functools.partial(jax.jit, static_argnums=(1, 3))
def g(a, b, c):
return a + b + c

print(g(1., 2., 3.))

I propose that a warning is raised if an arg num or arg name doesn't exist such that code will still run if the arg num or name doesn't exist but the programmer is made aware of the issue (and can silence if necessary)

Is there a good reason for this not existing already?

[x] I can't find a previous issue on this

Issue 2:
Hi There,

I've just done an install of jax on a fresh install of Ubuntu 18.04 LTS --- as vanilla as possible.

I followed the installation instructions for CPU+GPU support.

When I try to import jax in python I get the following:

Python 3.6.9 (default, Nov 7 2019, 10:44:02)

Type "copyright", "credits" or "license" for more information.

IPython 5.5.0 -- An enhanced Interactive Python.

? -> Introduction and overview of IPython's features.

%quickref -> Quick reference.

help -> Python's own help system.

object? -> Details about 'object', use 'object??' for extra details.

In [1]: import jax

RuntimeError Traceback (most recent call last)

RuntimeError: module compiled against API version 0xc but this version of numpy is 0xb

RuntimeError Traceback (most recent call last)

RuntimeError: module compiled against API version 0xc but this version of numpy is 0xb

ImportError Traceback (most recent call last)

in ()

----> 1 import jax

/home/case/.local/lib/python3.6/site-packages/jax/init.py in ()

17

18 from jax.version import version

---> 19 from jax.api import *

20 from jax import nn

21 from jax import random

/home/case/.local/lib/python3.6/site-packages/jax/api.py in ()

35 from contextlib import contextmanager

36

---> 37 from . import core

38 from . import linear_util as lu

39 from . import ad_util

/home/case/.local/lib/python3.6/site-packages/jax/core.py in ()

28 import numpy as onp

29

---> 30 from . import dtypes

31 from . import linear_util as lu

32

/home/case/.local/lib/python3.6/site-packages/jax/dtypes.py in ()

29 from . import util

30 from .config import flags

---> 31 from .lib import xla_client

32

33 FLAGS = flags.FLAGS

/home/case/.local/lib/python3.6/site-packages/jax/lib/init.py in ()

50 except:

51 tpu_client = None

---> 52 from jaxlib import xla_client

53 from jaxlib import lapack

54

/home/case/.local/lib/python3.6/site-packages/jaxlib/xla_client.py in ()

35 # and TensorFlow may fail with duplicate protocol buffer message definitions.

36

---> 37 from . import xla_extension as _xla

38 from .xla_extension import ops

39

ImportError: Unable to initialize Numpy API

I've tried searching for this issue and looking through these archives, but can't find anyone with this problem.

Any idea where I should start with trying to figure out the install?

As mentioned, I tried to keep everything as vanilla as possible: fresh Ubuntu 18.04 install, installed nvidia drivers with no issue, installed cuda 10.1 with no issue, the install instructions appeared to worked for jax.

So it seems quite possible that others may run into the issue.

Any help appreciated.