Regression in performance of samplers
With jax versions > 0.1.24, random samplers are slow. I think that it is due to recent changes on how jit works with static_args/kwargs. A script to reproduce,
```
import time
from jax import random

t = time.time()
random.normal(random.PRNGKey(0), shape=(1,))
print(time.time() - t)
t = time.time()
random.normal(random.PRNGKey(0), shape=(1,))
print(time.time() - t)
```
which returns `0.12923526763916016` and `0.12221717834472656`.

However, if we wrap these samplers in some function, then it is fast. For example,
```
def f():
    return random.normal(random.PRNGKey(0), shape=(1,))

t = time.time()
f()
print(time.time() - t)
t = time.time()
f()
print(time.time() - t)
```
will return `0.12787413597106934` and `0.0010831356048583984`.

I think that there is a small bug elsewhere which forces the sampler recompile. If this is an expected behaviour, then which function we should use to wrap these samplers to make it not recompiled?
cc @neerajprad

diff --git a/jax/api.py b/jax/api.py
index 7d7ae951d..5ce6f909d 100644
--- a/jax/api.py
+++ b/jax/api.py
@@ -46,7 +46,7 @@ from .tree_util import (process_pytree, node_types, build_tree, PyTreeDef,
                         tree_map, tree_flatten, tree_unflatten, tree_structure,
                         tree_transpose, leaf)
 from .util import (unzip2, unzip3, curry, partial, safe_map, safe_zip,
-                   WrapHashably, prod)
+                   WrapHashably, Hashable, prod)
 from .lib.xla_bridge import canonicalize_dtype, device_count
 from .abstract_arrays import ShapedArray
 from .interpreters import partial_eval as pe
@@ -824,11 +824,19 @@ def _argnums_partial(f, dyn_argnums, args):
     dyn_argnums = (dyn_argnums,)
   else:
     dyn_argnums = tuple(dyn_argnums)
-  fixed_args = tuple([None if i in dyn_argnums else WrapHashably(arg)
+  fixed_args = tuple([None if i in dyn_argnums else _wrap_hashably(arg)
                       for i, arg in enumerate(args)])
   dyn_args = tuple(args[i] for i in dyn_argnums)
   return _argnums_partial_(f, dyn_argnums, fixed_args), dyn_args
 
+def _wrap_hashably(arg):
+  try:
+    hash(arg)
+  except TypeError:
+    return WrapHashably(arg)
+  else:
+    return Hashable(arg)
+
 @lu.transformation
 def _argnums_partial_(dyn_argnums, fixed_args, *dyn_args, **kwargs):
   args = [None if arg is None else arg.val for arg in fixed_args]
diff --git a/jax/util.py b/jax/util.py
index 8faea1361..f54cb022d 100644
--- a/jax/util.py
+++ b/jax/util.py
@@ -189,6 +189,8 @@ def prod(xs):
 
 
 class WrapHashably(object):
+  __slots__ = ["val"]
+
   def __init__(self, val):
     self.val = val
 
@@ -198,6 +200,18 @@ class WrapHashably(object):
   def __eq__(self, other):
     return self.val is other.val
 
+class Hashable(object):
+  __slots__ = ["val"]
+
+  def __init__(self, val):
+    self.val = val
+
+  def __hash__(self):
+    return hash(self.val)
+
+  def __eq__(self, other):
+    return self.val == other.val
+
 
 
 def get_module_functions(module):
['jax/util.py', 'jax/api.py']