Repository: google/jax

Issue:
Memory explosion when jitting dynamic_slice
The [following example][1] explodes in memory usage when calling the jitted functions. The culprit seems to be the dynamic_slice. It happens calling with or without the GPU.

Incidentally, it is also rather slow (twice slower than the plain Python version), so I welcome comments if there is a jax-friendlier way to rewrite it, but perhaps that is what jit should do?

My system:

Fedora 29, 31 GB RAM
jax 0.1.36
jaxlib 0.1.16, installed from the googleapis wheel.
CUDA 10.0
Nvidia driver 430.14.

   [1]: https://gist.github.com/Dapid/b33919674f6a3cd9839473bfd2821e28 


Discussion:
Sorry for the long delay in response, most of us were away last week.

I think the memory blow-up happens because the gradient of `dynamic_slice` constructs a dense array of zeros to scatter a value into, and we end up with one such constant for each `dynamic_slice` call. This is similar to #850 ; in general JAX probably needs to get smarter about not instantiating large constants.

However, there's a better way to write your computation to use a single gather (via NumPy advanced indexing) rather than a long chain of `dynamic_slice` calls. The JIT-ted version no longer blows up in memory and is much faster:

```
    def single_loss(self, datum):
        J = self.J
        N = self.N

        loss = 0.
        for r in range(N):
            x = np.sum(J[np.arange(N), r, datum[np.arange(N)], datum[r]])
            loss -= np.sum(np.exp(x))
        return loss
```

If you want to get even fancier you can use a single gather for the entire outer loop too:

```
        xs = np.arange(N)
        ys = np.sum(J[xs[:, None], xs, datum[xs, None], datum[xs]], axis=1)
        loss = -np.sum(np.exp(ys))
```

(The second version didn't seem faster to me but it's fun to code-golf these things!)


Resolving patch:
diff --git a/jax/core.py b/jax/core.py
--- a/jax/core.py
+++ b/jax/core.py
@@ -119,7 +119,10 @@ def __eq__(self, other):
     return self.val == other.val if self.hashable else self.val is other.val
 
   def __repr__(self):
-    return 'Literal(val={}, hashable={})'.format(self.val, self.hashable)
+    if self.hashable:
+      return '{}'.format(self.val)
+    else:
+      return 'Literal(val={}, hashable={})'.format(self.val, self.hashable)
 
 class Primitive(object):
   def __init__(self, name):
diff --git a/jax/interpreters/partial_eval.py b/jax/interpreters/partial_eval.py
--- a/jax/interpreters/partial_eval.py
+++ b/jax/interpreters/partial_eval.py
@@ -454,6 +454,7 @@ def tracers_to_jaxpr(in_tracers, out_tracer):
   eqns = []
   env = {}
   consts = {}
+  const_to_var = defaultdict(newvar)
   destructuring_vars = {}
   for t in sorted_tracers:
     recipe = t.recipe
@@ -464,7 +465,8 @@ def tracers_to_jaxpr(in_tracers, out_tracer):
     elif isinstance(recipe, FreeVar):
       env[var(t)] = recipe.val
     elif isinstance(recipe, ConstVar):
-      consts[var(t)] = recipe.val
+      v = t_to_var[id(t)] = const_to_var[id(recipe.val)]
+      consts[v] = recipe.val
     elif isinstance(recipe, Literal):
       t_to_var[id(t)] = recipe
     elif isinstance(recipe, Destructuring):


Test patch: