Repository: jax-ml/jax
Pull request: Convert int input to static_argnums to a tuple
User could make mistake of passing an int to static_argnums, this helps to avoid unnecessary error.

Fix #969 

Problem statement:
convert integer for `static_argnums` to tuple automatically
I've made the same mistake a few times on `static_argnums` on `jit` function, by not using a tuple 

e.g., `jit(f, static_argnums=1)`
would it be possible / reasonable to add a simple `if` statement to convert it to tuple?

where the correct way is  `jit(f, static_argnums=(1,) )`

also note `(1)` is not a tuple,  but `(1,)` is

Hints from issue discussion:

Commit message: Convert int input to static_argnums to a tuple

User could make mistake of passing an int to static_argnums, this helps to avoid unnecessary error.

Patch to review:
diff --git a/jax/api.py b/jax/api.py
index 047df3c9f4c5..b36acc1bf417 100644
--- a/jax/api.py
+++ b/jax/api.py
@@ -110,6 +110,9 @@ def _jit(fun, static_argnums, device_values=True):
   def f_jitted(*args, **kwargs):
     if _jit_is_disabled or config.read('jax_disable_jit'):
       return fun(*args, **kwargs)
+    if isinstance(static_argnums, int):
+      # static_argnums is a tuple of ints
+      static_argnums = tuple([static_argnums])
     if static_argnums and max(static_argnums) >= len(args):
       msg = ("Jitted function has static_argnums={} but was called with only {}"
              " positional arguments.")


Code review comments:
diff_hunk: @@ -110,6 +110,9 @@ def _jit(fun, static_argnums, device_values=True):
   def f_jitted(*args, **kwargs):
     if _jit_is_disabled or config.read('jax_disable_jit'):
       return fun(*args, **kwargs)
+    if isinstance(static_argnums, int):
+      # static_argnums is a tuple of ints
+      static_argnums = tuple([static_argnums])
path: jax/api.py
text: @user1:
How about `static_argnums = (static_argnums,)`?