<p dir="auto">Currently, if a static_argnums or static_argname doesn't exist for the function, nothing happens<br>
This can be an issue if a function is updated but the jax.jit not being updated</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="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.))"><pre class="notranslate"><code class="notranslate">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.))
</code></pre></div>
<p dir="auto">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)<br>
Is there a good reason for this not existing already?</p>
<p dir="auto">[x] I can't find a previous issue on this</p>
<p dir="auto">Hi There,</p>
<p dir="auto">I've just done an install of jax on a fresh install of Ubuntu 18.04 LTS --- as vanilla as possible.<br>
I followed the installation instructions for CPU+GPU support.</p>
<p dir="auto">When I try to import jax in python I get the following:</p>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
<p dir="auto">Python 3.6.9 (default, Nov  7 2019, 10:44:02)<br>
Type "copyright", "credits" or "license" for more information.</p>
<p dir="auto">IPython 5.5.0 -- An enhanced Interactive Python.<br>
?         -&gt; Introduction and overview of IPython's features.<br>
%quickref -&gt; Quick reference.<br>
help      -&gt; Python's own help system.<br>
object?   -&gt; Details about 'object', use 'object??' for extra details.</p>
<h2 dir="auto">In [1]: import jax</h2>
<h2 dir="auto">RuntimeError                              Traceback (most recent call last)<br>
RuntimeError: module compiled against API version 0xc but this version of numpy is 0xb</h2>
<h2 dir="auto">RuntimeError                              Traceback (most recent call last)<br>
RuntimeError: module compiled against API version 0xc but this version of numpy is 0xb</h2>
<p dir="auto">ImportError                               Traceback (most recent call last)<br>
 in ()<br>
----&gt; 1 import jax</p>
<p dir="auto">/home/case/.local/lib/python3.6/site-packages/jax/<strong>init</strong>.py in ()<br>
17<br>
18 from jax.version import <strong>version</strong><br>
---&gt; 19 from jax.api import *<br>
20 from jax import nn<br>
21 from jax import random</p>
<p dir="auto">/home/case/.local/lib/python3.6/site-packages/jax/api.py in ()<br>
35 from contextlib import contextmanager<br>
36<br>
---&gt; 37 from . import core<br>
38 from . import linear_util as lu<br>
39 from . import ad_util</p>
<p dir="auto">/home/case/.local/lib/python3.6/site-packages/jax/core.py in ()<br>
28 import numpy as onp<br>
29<br>
---&gt; 30 from . import dtypes<br>
31 from . import linear_util as lu<br>
32</p>
<p dir="auto">/home/case/.local/lib/python3.6/site-packages/jax/dtypes.py in ()<br>
29 from . import util<br>
30 from .config import flags<br>
---&gt; 31 from .lib import xla_client<br>
32<br>
33 FLAGS = flags.FLAGS</p>
<p dir="auto">/home/case/.local/lib/python3.6/site-packages/jax/lib/<strong>init</strong>.py in ()<br>
50 except:<br>
51   tpu_client = None<br>
---&gt; 52 from jaxlib import xla_client<br>
53 from jaxlib import lapack<br>
54</p>
<p dir="auto">/home/case/.local/lib/python3.6/site-packages/jaxlib/xla_client.py in ()<br>
35 # and TensorFlow may fail with duplicate protocol buffer message definitions.<br>
36<br>
---&gt; 37 from . import xla_extension as _xla<br>
38 from .xla_extension import ops<br>
39</p>
<p dir="auto">ImportError: Unable to initialize Numpy API</p>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
<p dir="auto">I've tried searching for this issue and looking through these archives, but can't find anyone with this problem.<br>
Any idea where I should start with trying to figure out the install?<br>
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.<br>
So it seems quite possible that others may run into the issue.</p>
<p dir="auto">Any help appreciated.</p>