Error differentiating QR decomposition
This appears to be a bug but unfortunately I don't have time to delve into it.

Below is a simple script to do gradient descent on a function involving a QR decomposition. It fails when I execute `training_gradient_fun` with `TypeError: No abstraction handler for type: <class 'jax.ad_util.Zero'>`.

Ubuntu 16.04
python 3.6.6
jax 0.1.20
jaxlib 0.1.9

```
import jax
import jax.numpy as np
import numpy as onp

def loss(weights, inputs):
    q, r = np.linalg.qr(weights)
    a = np.matmul(q, weights)
    return np.mean(a ** 2)

training_gradient_fun = jax.jit(jax.grad(loss))

inputs = onp.random.rand(5, 5)
weights = onp.random.rand(5, 5)

for i in range(100):
    weights -= 0.1 * training_gradient_fun(weights, inputs)
```
diff --git a/jax/lax.py b/jax/lax.py
index 3ffc2598f..556886e99 100644
--- a/jax/lax.py
+++ b/jax/lax.py
@@ -1363,7 +1363,7 @@ ad.primitive_transposes[add_p] = _add_transpose
 
 def _sub_transpose(t, x, y):
   assert x is None and y is None  # computation must be linear, not affine
-  return [t, neg(t)]
+  return [t, neg(t) if t is not ad_util.zero else ad_util.zero]
 
 sub_p = standard_binop([_num, _num], 'sub')
 ad.defjvp(sub_p,
['jax/lax.py']