Repository: google/jax
Issue: make pmap compilation cache "act globally" with no axis name

Both the `jit` and `pmap` compilation caches are global. The `jit` cache acts that way:

```python
from jax import jit

def f(x):
  print("miss")
  return x

jit(f)(3)  # prints
jit(f)(3)  # doesn't print, i.e. cache hit
```

But the `pmap` one seems not to when you don't provide a name for the axis:

```python
from jax import pmap
from jax.lib import xla_bridge as xb
import jax.numpy as np

def f(x):
  print("miss")
  return x

x = np.ones(xb.device_count())
pmap(f)(x)  # prints
pmap(f)(x)  # prints

pmap(f, axis_name='i')(x)  # prints
pmap(f, axis_name='i')(x)  # doesn't print
```

The issue is that when an axis name isn't provided, we create a dummy one, and we create a new one each time we evaluate `pmap(f)` to produce a new pmapped version of `f`. We should probably remedy that, since patterns like

```python
for i in range(100):
  x = pmap(update)(x)
```

are handy, but they're currently a foot-gun with respect to compilation caching.

Changed files:
jax/linear_util.py, jax/interpreters/ad.py, jax/api.py, tests/api_test.py

diff --git a/jax/api.py b/jax/api.py
index e742fbab8..6dd64aef6 100644
--- a/jax/api.py
+++ b/jax/api.py
@@ -846,7 +846,7 @@ def pmap(fun, axis_name=None, devices=None, backend=None):
   [ 13.  13.]
   """
   _check_callable(fun)
-  axis_name = _TempAxisName() if axis_name is None else axis_name
+  axis_name = _TempAxisName(fun) if axis_name is None else axis_name
 
   @wraps(fun)
   def f_pmapped(*args, **kwargs):
@@ -875,13 +875,19 @@ def _pmap_axis_size(xs):
     raise ValueError(msg.format([x for x in xs if not hasattr(x, 'shape')]))
 
 class _TempAxisName(object):
+  def __init__(self, obj):
+    self.obj = obj
   def __repr__(self):
-    return '<axis {}>'.format(hex(id(self)))
+    return '<axis {}>'.format(hex(id(self.obj)))
+  def __hash__(self):
+    return hash(self.obj)
+  def __eq__(self, other):
+    return self.obj is other.obj
 
 
 def soft_pmap(fun, axis_name=None, backend=None):
   _check_callable(fun)
-  axis_name = _TempAxisName() if axis_name is None else axis_name
+  axis_name = _TempAxisName(fun) if axis_name is None else axis_name
 
   @wraps(fun)
   def f_pmapped(*args, **kwargs):
@@ -930,7 +936,7 @@ def _reshape_merge(x):
 
 def _papply(fun):
   # This function is for testing purposes.
-  axis_name = _TempAxisName()
+  axis_name = _TempAxisName(fun)
 
   def papply_fun(*args, **kwargs):
     f = lu.wrap_init(fun)
@@ -944,7 +950,7 @@ def _papply(fun):
 
 
 def _parallelize(fun):
-  axis_name = _TempAxisName()
+  axis_name = _TempAxisName(fun)
 
   def pfun(*args):
     f = lu.wrap_init(fun)
diff --git a/jax/interpreters/ad.py b/jax/interpreters/ad.py
index 36e0cfa61..3b2f42424 100644
--- a/jax/interpreters/ad.py
+++ b/jax/interpreters/ad.py
@@ -26,7 +26,7 @@ from ..ad_util import (add_jaxvals, add_jaxvals_p, zeros_like_jaxval, zeros_like
 from ..abstract_arrays import raise_to_shaped
 from ..util import unzip2, unzip3, safe_map, safe_zip, partial, split_list
 from ..tree_util import build_tree, register_pytree_node, tree_map
-from ..linear_util import thunk, staged, transformation, transformation_with_aux, wrap_init
+from ..linear_util import thunk, transformation, transformation_with_aux, wrap_init
 from ..api_util import flatten_fun, flatten_fun_nokwargs
 from ..tree_util import tree_flatten, tree_unflatten
 
diff --git a/jax/linear_util.py b/jax/linear_util.py
index 59d408614..b71305497 100644
--- a/jax/linear_util.py
+++ b/jax/linear_util.py
@@ -69,7 +69,7 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-import fastcache
+import weakref
 
 from .util import curry, partial
 
@@ -112,18 +112,6 @@ class Store(object):
   __bool__ = __nonzero__
 
 
-@curry
-def staged(f, *init_args):
-  store = Store()
-  def f_partial(*rest):
-    ans, aux = f(*(init_args + rest))
-    store.store(aux)
-    return ans
-
-  f_partial.__name__ = f.__name__ + "_staged"
-  return f_partial, thunk(lambda: store.val)
-
-
 class WrappedFun(object):
   """Represents a function `f` to which `transforms` are to be applied.
 
@@ -149,8 +137,8 @@ class WrappedFun(object):
     return WrappedFun(self.f, ((gen, gen_args),) + self.transforms,
                       (out_store,) + self.stores, self.params)
 
-  def populate_stores(self, other):
-    for self_store, other_store in zip(self.stores, other.stores):
+  def populate_stores(self, stores):
+    for self_store, other_store in zip(self.stores, stores):
       if self_store is not None:
         self_store.store(other_store.val)
 
@@ -208,15 +196,17 @@ def wrap_init(f, params={}):
   return WrappedFun(f, (), (), tuple(sorted(params.items())))
 
 
-def cache(call, max_size=4096):
-  @fastcache.clru_cache(maxsize=max_size)
-  def cached_fun_body(f, args):
-    return call(f, *args), f
-
-  def cached_fun(f, *args):
-    ans, f_prev = cached_fun_body(f, args)
-    if id(f_prev) != id(f):
-      f.populate_stores(f_prev)
+def cache(call):
+  fun_caches = weakref.WeakKeyDictionary()
+  def memoized_fun(fun, *args):
+    cache = fun_caches.setdefault(fun.f, {})
+    key = (fun.transforms, fun.params, args)
+    result = cache.get(key, None)
+    if result is not None:
+      ans, stores = result
+      fun.populate_stores(stores)
+    else:
+      ans = call(fun, *args)
+      cache[key] = (ans, fun.stores)
     return ans
-
-  return cached_fun
+  return memoized_fun
diff --git a/tests/api_test.py b/tests/api_test.py
index e429be63a..071faa58f 100644
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -20,6 +20,7 @@ import collections
 from functools import partial
 import unittest
 import warnings
+import weakref
 
 from absl.testing import absltest
 import numpy as onp
@@ -1181,6 +1182,48 @@ class APITest(jtu.JaxTestCase):
         foo, in_axes=((0, collections.OrderedDict([('a', 1), ('b', 2)])),))
     self.assertEqual(vfoo(tree).shape, (6, 2, 5))
 
+  def test_jit_reference_dropping(self):
+    x = onp.ones(10)
+    f = (lambda x: lambda: x)(x)  # reference to x in f's closure
+    g = jit(f)
+    x = weakref.ref(x)      # no more strong ref to x in this scope
+    assert x() is not None  # x is still around
+    f()                     # f runs
+    g()                     # g runs
+    g()                     # g runs a second time
+    del f                   # delete the raw callable
+    assert x() is not None  # x is still around
+    g()                     # g still runs
+    del g                   # no more references to x
+    assert x() is None      # x is gone
+
+  def test_jit_global_cache(self):
+    def f(x):
+      assert python_should_be_executing
+      return x
+
+    python_should_be_executing = True
+    api.jit(f)(2)
+    python_should_be_executing = False
+    api.jit(f)(3)
+
+  def test_pmap_global_cache(self):
+    def f(x):
+      assert python_should_be_executing
+      return x
+
+    x = onp.ones(1)
+
+    python_should_be_executing = True
+    api.pmap(f)(x)
+    python_should_be_executing = False
+    api.pmap(f)(x)
+
+    python_should_be_executing = True
+    api.pmap(f, 'i')(x)
+    python_should_be_executing = False
+    api.pmap(f, 'i')(x)
+
 
 if __name__ == '__main__':
   absltest.main()