Repository: google/jax
Issue: grad(vmap(cholesky)) error

I'm getting an error when trying to call grad(vmap(cholesky)):

```python
psd_mat = onp.random.randn(20, 10)
psd_mat = psd_mat.T.dot(psd_mat)
vec = onp.random.randn(10)

def f(scale):
  scaled_mat = scale * psd_mat
  chol = np.linalg.cholesky(scaled_mat)
  return -0.5 * np.sum((np.einsum('ij,j->i', chol, vec))**2)
vmapped_f = jax.vmap(f)
vmapped_f_grad = jax.grad(vmapped_f)

# Works fine
print(f(0.1))

# Also works fine
print(vmapped_f(0.1 + onp.zeros((5, 1))))

# Throws exception
print(vmapped_f_grad(0.1 + onp.zeros((5, 1))))
```

Full output is
```
/usr/local/lib/python2.7/dist-packages/jax/numpy/linalg.py:53: UserWarning: numpy.linalg support is experimental and may cause silent failures or wrong outputs
  warnings.warn(_EXPERIMENTAL_WARNING)
/usr/local/lib/python2.7/dist-packages/jax/lib/xla_bridge.py:146: UserWarning: No GPU found, falling back to CPU.
  warnings.warn('No GPU found, falling back to CPU.')
-8.891344
[-8.891344 -8.891344 -8.891344 -8.891344 -8.891344]

TypeErrorTraceback (most recent call last)
<ipython-input-3-73aa1b00356c> in <module>()
     14 
     15 vmapped_f_grad = jax.grad(vmapped_f)
---> 16 print(vmapped_f_grad(0.1 + onp.zeros((5, 1))))

/usr/local/lib/python2.7/dist-packages/jax/api.pyc in grad_f(*args, **kwargs)
    112   @wraps(fun, docstr=docstr, argnums=argnums)
    113   def grad_f(*args, **kwargs):
--> 114     ans, g = value_and_grad_f(*args, **kwargs)
    115     return g
    116 

/usr/local/lib/python2.7/dist-packages/jax/api.pyc in value_and_grad_f(*args, **kwargs)
    147     f = lu.wrap_init(fun, kwargs)
    148     f_partial, dyn_args = argnums_partial(f, argnums, args)
--> 149     ans, vjp_py = vjp(f_partial, *dyn_args)
    150     check_scalar(ans)
    151     g = vjp_py(onp.ones((), onp.result_type(ans)))

/usr/local/lib/python2.7/dist-packages/jax/api.pyc in vjp(fun, *primals)
    358   check_args(primals_flat)
    359   jaxtree_fun, out_tree = pytree_fun_to_jaxtupletree_fun(fun, in_trees)
--> 360   out_primal, out_vjp = ad.vjp(jaxtree_fun, primals_flat)
    361   out_tree = out_tree()
    362   out_primal_py = build_tree(out_tree, out_primal)

/usr/local/lib/python2.7/dist-packages/jax/interpreters/ad.pyc in vjp(traceable, primals)
     72 
     73 def vjp(traceable, primals):
---> 74   out_primal, pval, jaxpr, consts = linearize(traceable, *primals)
     75   def vjp_(ct):
     76     ct = ignore_consts(ct, pval)

/usr/local/lib/python2.7/dist-packages/jax/interpreters/ad.pyc in linearize(traceable, *primals)
     65   in_pvals = (pe.PartialVal((None, pack(primals))),
     66               pe.PartialVal((core.AbstractTuple(tangent_avals), core.unit)))
---> 67   jaxpr, out_pval, consts = pe.trace_to_jaxpr(jvpfun, in_pvals)
     68   pval_primal, pval_tangent = unpair_pval(out_pval)
     69   aval_primal, const_primal = pval_primal

/usr/local/lib/python2.7/dist-packages/jax/interpreters/partial_eval.pyc in trace_to_jaxpr(fun, pvals, **kwargs)
    254   with new_master(JaxprTrace) as master:
    255     fun = trace_to_subjaxpr(fun, master)
--> 256     jaxpr, (out_pval, consts, env) = fun.call_wrapped(pvals, **kwargs)
    257     assert not env
    258     del master

/usr/local/lib/python2.7/dist-packages/jax/linear_util.pyc in call_wrapped(self, *args)
     84 
     85     del gen
---> 86     ans = self.f(*args, **self.kwargs)
     87     del args
     88     while stack:

/usr/local/lib/python2.7/dist-packages/jax/api.pyc in batched_fun(*args, **kwargs)
    253     in_flat, in_trees = unzip2(map(pytree_to_jaxtupletree, args))
    254     jaxtree_fun, out_tree = pytree_fun_to_jaxtupletree_fun(f, in_trees)
--> 255     out_flat = batching.batch(jaxtree_fun, in_flat, in_axes_, out_axes)
    256     return build_tree(out_tree(), out_flat)
    257 

/usr/local/lib/python2.7/dist-packages/jax/interpreters/batching.pyc in batch(fun, in_vals, in_dims, out_dim_target)
     41     return fun.call_wrapped(*in_vals), None  # no mapped dimensions
     42   elif len(sizes) == 1:
---> 43     out_val, out_dim = batch_transform(fun).call_wrapped(in_vals, in_dims)
     44     return moveaxis(sizes.pop(), out_dim_target, out_dim, out_val)
     45   else:

/usr/local/lib/python2.7/dist-packages/jax/linear_util.pyc in call_wrapped(self, *args)
     84 
     85     del gen
---> 86     ans = self.f(*args, **self.kwargs)
     87     del args
     88     while stack:

<ipython-input-3-73aa1b00356c> in f(scale)
      5 def f(scale):
      6   scaled_mat = scale * psd_mat
----> 7   chol = np.linalg.cholesky(scaled_mat)
      8   return -0.5 * np.sum((np.dot(chol, vec))**2)
      9 

/usr/local/lib/python2.7/dist-packages/jax/numpy/linalg.pyc in cholesky(a)
     53   warnings.warn(_EXPERIMENTAL_WARNING)
     54   a = _promote_arg_dtypes(np.asarray(a))
---> 55   return lax_linalg.cholesky(a)
     56 
     57 

/usr/local/lib/python2.7/dist-packages/jax/lax_linalg.pyc in cholesky(x)
     35 # traceables
     36 
---> 37 def cholesky(x): return cholesky_p.bind(x)
     38 
     39 def eigh(x, lower=True): return eigh_p.bind(x, lower=lower)

/usr/local/lib/python2.7/dist-packages/jax/core.pyc in bind(self, *args, **kwargs)
     72 
     73     tracers = map(top_trace.full_raise, args)
---> 74     out_tracer = top_trace.process_primitive(self, tracers, kwargs)
     75     return full_lower(out_tracer)
     76 

/usr/local/lib/python2.7/dist-packages/jax/interpreters/batching.pyc in process_primitive(self, primitive, tracers, params)
    120       # TODO(mattjj,phawkins): if no rule implemented, could vmap-via-map here
    121       batched_primitive = get_primitive_batcher(primitive)
--> 122       val_out, dim_out = batched_primitive(vals_in, dims_in, **params)
    123       return BatchTracer(self, val_out, dim_out)
    124 

/usr/local/lib/python2.7/dist-packages/jax/lax_linalg.pyc in cholesky_batching_rule(batched_args, batch_dims)
     89   bd, = batch_dims
     90   x = batching.bdim_at_front(x, bd)
---> 91   return cholesky(x), 0
     92 
     93 cholesky_p = standard_unop(_float | _complex, 'cholesky')

/usr/local/lib/python2.7/dist-packages/jax/lax_linalg.pyc in cholesky(x)
     35 # traceables
     36 
---> 37 def cholesky(x): return cholesky_p.bind(x)
     38 
     39 def eigh(x, lower=True): return eigh_p.bind(x, lower=lower)

/usr/local/lib/python2.7/dist-packages/jax/core.pyc in bind(self, *args, **kwargs)
     72 
     73     tracers = map(top_trace.full_raise, args)
---> 74     out_tracer = top_trace.process_primitive(self, tracers, kwargs)
     75     return full_lower(out_tracer)
     76 

/usr/local/lib/python2.7/dist-packages/jax/interpreters/ad.pyc in process_primitive(self, primitive, tracers, params)
    178           "Forward-mode differentiation rule for '{}' not implemented"
    179           .format(primitive))
--> 180     primal_out, tangent_out = jvp(primals_in, tangents_in, **params)
    181     return JVPTracer(self, primal_out, tangent_out)
    182 

/usr/local/lib/python2.7/dist-packages/jax/lax_linalg.pyc in cholesky_jvp_rule(primals, tangents)
     82                          left_side=False, transpose_a=True, lower=True)
     83   L_dot = lax.dot(L, phi(triangular_solve(
---> 84       L, tmp, left_side=True, transpose_a=False, lower=True)))
     85   return L, L_dot
     86 

/usr/local/lib/python2.7/dist-packages/jax/lax.pyc in dot(lhs, rhs)
    190       rhs_shape=rhs.shape)
    191 
--> 192 def dot(lhs, rhs): return dot_p.bind(lhs, rhs)
    193 
    194 def dot_general(lhs, rhs, dimension_numbers):

/usr/local/lib/python2.7/dist-packages/jax/core.pyc in bind(self, *args, **kwargs)
     72 
     73     tracers = map(top_trace.full_raise, args)
---> 74     out_tracer = top_trace.process_primitive(self, tracers, kwargs)
     75     return full_lower(out_tracer)
     76 

/usr/local/lib/python2.7/dist-packages/jax/interpreters/partial_eval.pyc in process_primitive(self, primitive, tracers, params)
     67     tracers = map(self.instantiate_const, tracers)
     68     avals = [t.aval for t in tracers]
---> 69     out_aval = primitive.abstract_eval(*avals, **params)
     70     eqn = JaxprEqn(tracers, None, primitive, (), False, params)
     71     return JaxprTracer(self, PartialVal((out_aval, unit)), eqn)

/usr/local/lib/python2.7/dist-packages/jax/lax.pyc in standard_abstract_eval(shape_rule, dtype_rule, *args, **kwargs)
    753     return ShapedArray(shape_rule(*args, **kwargs), dtype_rule(*args, **kwargs))
    754   elif least_specialized is ShapedArray:
--> 755     return ShapedArray(shape_rule(*args, **kwargs), dtype_rule(*args, **kwargs))
    756   elif least_specialized is UnshapedArray:
    757     return UnshapedArray(dtype_rule(*args, **kwargs))

/usr/local/lib/python2.7/dist-packages/jax/lax.pyc in _dot_shape_rule(lhs, rhs)
   1227   if lhs.ndim > 2 or rhs.ndim > 2:
   1228     msg = "Dot only supports rank 2 or less, got shapes {} and {}."
-> 1229     raise TypeError(msg.format(lhs.shape, rhs.shape))
   1230 
   1231   def require(shape_cond):

TypeError: Dot only supports rank 2 or less, got shapes (5, 10, 10) and (5, 10, 10).
```

Changed files:
jax/numpy/lax_numpy.py, jax/scipy/linalg.py, tests/batching_test.py, tests/linalg_test.py, jax/lax.py, tests/lax_numpy_test.py, jax/numpy/linalg.py, jax/lax_linalg.py

diff --git a/jax/lax.py b/jax/lax.py
index 6d0c030ff..9511f1f01 100644
--- a/jax/lax.py
+++ b/jax/lax.py
@@ -831,10 +831,12 @@ standard_binop = partial(binop, _input_dtype)
 # a broadcast). but saving the shape info with the primitives isn't great either
 # because then we can't trace these ops without shape data.
 def _brcast(x, *others):
-  # used in jvprules to make binop broadcasting explicit for transposability.
-  # requires shape info during jvp tracing, which isn't strictly necessary.
-  shapes = list(filter(None, map(onp.shape, (x,) + others)))
-  shape = tuple(shapes and onp.max(shapes, axis=0))
+  # Used in jvprules to make binop broadcasting explicit for transposability.
+  # Requires shape info during jvp tracing, which isn't strictly necessary.
+  # We don't need full numpy broadcasting, but otherwise the logic is the same
+  # so we reuse the broadcast_shapes function after filtering out scalars.
+  shapes = tuple(filter(None, map(onp.shape, (x,) + others)))
+  shape = shapes and broadcast_shapes(*shapes)
   if onp.shape(x) != shape:
     return _brcast_to(x, shape)
   else:
@@ -1355,10 +1357,6 @@ def _dot_general_shape_rule(lhs, rhs, dimension_numbers):
     msg = ("dot_general requires rhs batch dimensions to precede contracting "
            "and non-contracting dimensions, got rhs_batch {}.")
     raise TypeError(msg.format(rhs_batch))
-  if not len(lhs_contracting) == len(rhs_contracting) == 1:
-    msg = ("dot_general accepts exactly one lhs_contracting and "
-           "rhs_contracting dimension, got {} and {}.")
-    raise TypeError(msg.format(lhs_contracting, rhs_contracting))
   lhs_contracting_shape = onp.take(lhs.shape, lhs_contracting)
   rhs_contracting_shape = onp.take(rhs.shape, rhs_contracting)
   if not onp.all(onp.equal(lhs_contracting_shape, rhs_contracting_shape)):
diff --git a/jax/lax_linalg.py b/jax/lax_linalg.py
index 10c118ca7..73d553880 100644
--- a/jax/lax_linalg.py
+++ b/jax/lax_linalg.py
@@ -34,7 +34,10 @@ from jaxlib import lapack
 
 # traceables
 
-def cholesky(x): return cholesky_p.bind(x)
+def cholesky(x, symmetrize_input=True):
+  if symmetrize_input:
+    x = symmetrize(x)
+  return cholesky_p.bind(x)
 
 def eigh(x, lower=True): return eigh_p.bind(x, lower=lower)
 
@@ -60,8 +63,9 @@ def triangular_solve(a, b, left_side=False, lower=False, transpose_a=False,
 
 # utilities
 
-def _T(x):
-  return np.swapaxes(x, -1, -2)
+def _T(x): return np.swapaxes(x, -1, -2)
+def _H(x): return np.conj(_T(x))
+def symmetrize(x): return (x + _H(x)) / 2
 
 
 # primitives
@@ -76,11 +80,10 @@ def cholesky_jvp_rule(primals, tangents):
   L = cholesky_p.bind(x)
 
   # Forward-mode rule from https://arxiv.org/pdf/1602.07527.pdf
-  sigma_dot = (sigma_dot + _T(sigma_dot)) / 2
-  phi = lambda X: np.tril(X) / (1 + np.eye(x.shape[-1]))
+  phi = lambda X: np.tril(X) / (1 + np.eye(X.shape[-1], dtype=X.dtype))
   tmp = triangular_solve(L, sigma_dot,
                          left_side=False, transpose_a=True, lower=True)
-  L_dot = lax.dot(L, phi(triangular_solve(
+  L_dot = lax.batch_matmul(L, phi(triangular_solve(
       L, tmp, left_side=True, transpose_a=False, lower=True)))
   return L, L_dot
 
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
index b16f133e4..81ecdb1ba 100644
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -326,8 +326,8 @@ logical_xor = _logical_op(onp.logical_xor, lax.bitwise_xor)
 
 @_wraps(onp.true_divide)
 def true_divide(x1, x2):
-  x1, x2 = _promote_shapes(x1, x2)
   result_dtype = _result_dtype(onp.true_divide, x1, x2)
+  x1, x2 = _promote_shapes(x1, x2)
   return lax.div(lax.convert_element_type(x1, result_dtype),
                  lax.convert_element_type(x2, result_dtype))
 
@@ -390,7 +390,7 @@ def _float_divmod(x1, x2):
 
 @_wraps(onp.logaddexp)
 def logaddexp(x1, x2):
-  x1, x2 = _promote_to_result_dtype(onp.logaddexp, *_promote_shapes(x1, x2))
+  x1, x2 = _promote_shapes(*_promote_to_result_dtype(onp.logaddexp, x1, x2))
   amax = lax.max(x1, x2)
   return lax.add(amax, lax.log(lax.add(lax.exp(lax.sub(x1, amax)),
                                        lax.exp(lax.sub(x2, amax)))))
@@ -398,7 +398,7 @@ def logaddexp(x1, x2):
 
 @_wraps(onp.logaddexp2)
 def logaddexp2(x1, x2):
-  x1, x2 = _promote_to_result_dtype(onp.logaddexp2, *_promote_shapes(x1, x2))
+  x1, x2 = _promote_shapes(*_promote_to_result_dtype(onp.logaddexp2, x1, x2))
   amax = lax.max(x1, x2)
   return lax.add(amax, log2(lax.add(exp2(lax.sub(x1, amax)),
                                     exp2(lax.sub(x2, amax)))))
diff --git a/jax/numpy/linalg.py b/jax/numpy/linalg.py
index ec086f673..888cc84d0 100644
--- a/jax/numpy/linalg.py
+++ b/jax/numpy/linalg.py
@@ -33,7 +33,6 @@ _EXPERIMENTAL_WARNING = "numpy.linalg support is experimental and may cause sile
 _T = lambda x: np.swapaxes(x, -1, -2)
 
 
-
 def _promote_arg_dtypes(*args):
   """Promotes `args` to a common inexact type."""
   def _to_inexact_type(type):
@@ -47,7 +46,6 @@ def _promote_arg_dtypes(*args):
     return args
 
 
-
 @_wraps(onp.linalg.cholesky)
 def cholesky(a):
   warnings.warn(_EXPERIMENTAL_WARNING)
diff --git a/jax/scipy/linalg.py b/jax/scipy/linalg.py
index 2077a0a5f..240984a6d 100644
--- a/jax/scipy/linalg.py
+++ b/jax/scipy/linalg.py
@@ -36,7 +36,7 @@ def cholesky(a, lower=False, overwrite_a=False, check_finite=True):
   warnings.warn(_EXPERIMENTAL_WARNING)
   del overwrite_a, check_finite
   a = np_linalg._promote_arg_dtypes(np.asarray(a))
-  l = lax_linalg.cholesky(a if lower else np.conj(_T(a)))
+  l = lax_linalg.cholesky(a if lower else np.conj(_T(a)), symmetrize_input=False)
   return l if lower else np.conj(_T(l))
 
 
diff --git a/tests/batching_test.py b/tests/batching_test.py
index b4eb21939..c573b57d0 100644
--- a/tests/batching_test.py
+++ b/tests/batching_test.py
@@ -793,6 +793,23 @@ class BatchingTest(jtu.JaxTestCase):
     expected = onp.transpose(x, (2, 1, 3, 0))
     self.assertAllClose(ans, expected, check_dtypes=False)
 
+  def testIssue354(self):
+    psd_mat = onp.random.randn(20, 10)
+    psd_mat = psd_mat.T.dot(psd_mat)
+    vec = onp.random.randn(10)
+
+    def f(scale):
+      scaled_mat = scale * psd_mat
+      chol = np.linalg.cholesky(scaled_mat)
+      return -0.5 * np.sum((np.einsum('ij,j->i', chol, vec))**2)
+    vmapped_f = vmap(f)
+    vmapped_f_grad = grad(lambda x: np.sum(vmapped_f(x)))
+
+    scales = onp.array([[0.1], [0.2], [0.3], [0.4], [0.5]])
+    ans = vmapped_f_grad(scales)  # don't crash!
+    expected = onp.stack([grad(f)(scale) for scale in scales])
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
 
 if __name__ == '__main__':
   absltest.main()
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
index 190c8710f..863a025f1 100644
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -1215,6 +1215,14 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
     jax_numpy_result = (1 + lnp.eye(1, dtype=lnp.float32)).dtype
     self.assertEqual(orig_numpy_result, jax_numpy_result)
 
+  def testSymmetrizeDtypePromotion(self):
+    x = onp.eye(3, dtype=onp.float32)
+    orig_numpy_result = ((x + x.T) / 2).dtype
+
+    x = lnp.eye(3, dtype=lnp.float32)
+    jax_numpy_result = ((x + x.T) / 2).dtype
+    self.assertEqual(orig_numpy_result, jax_numpy_result)
+
 
 if __name__ == "__main__":
   absltest.main()
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
index 3aea4517a..7eff490c7 100644
--- a/tests/linalg_test.py
+++ b/tests/linalg_test.py
@@ -61,13 +61,17 @@ class NumpyLinalgTest(jtu.JaxTestCase):
       for rng in [jtu.rand_default()]))
   def testCholesky(self, shape, dtype, rng):
     def args_maker():
-      a = rng(shape, dtype)
+      factor_shape = shape[:-1] + (2 * shape[-1],)
+      a = rng(factor_shape, dtype)
       return [onp.matmul(a, np.conj(T(a)))]
 
     self._CheckAgainstNumpy(onp.linalg.cholesky, np.linalg.cholesky, args_maker,
                             check_dtypes=True, tol=1e-3)
     self._CompileAndCheck(np.linalg.cholesky, args_maker, check_dtypes=True)
 
+    if onp.finfo(dtype).bits == 64:
+      jtu.check_grads(np.linalg.cholesky, args_maker(), order=2)
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name":
        "_n={}".format(jtu.format_shape_dtype_string((n,n), dtype)),
@@ -300,7 +304,6 @@ class NumpyLinalgTest(jtu.JaxTestCase):
                             check_dtypes=True, tol=1e-3)
     self._CompileAndCheck(np.linalg.solve, args_maker, check_dtypes=True)
 
-
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name":
        "_shape={}".format(jtu.format_shape_dtype_string(shape, dtype)),
@@ -374,7 +377,6 @@ class ScipyLinalgTest(jtu.JaxTestCase):
                             args_maker, check_dtypes=True, tol=1e-3)
     self._CompileAndCheck(jsp.linalg.lu_factor, args_maker, check_dtypes=True)
 
-
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name":
        "_lhs={}_rhs={}_sym_pos={}_lower={}".format(
@@ -414,7 +416,6 @@ class ScipyLinalgTest(jtu.JaxTestCase):
                             check_dtypes=True, tol=1e-3)
     self._CompileAndCheck(jsp_fun, args_maker, check_dtypes=True)
 
-
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name":
        "_lhs={}_rhs={}_lower={}_transposea={}".format(
@@ -482,5 +483,6 @@ class ScipyLinalgTest(jtu.JaxTestCase):
                 trans=1 if transpose_a else 0)
     jtu.check_grads(f, (A, B), 2, rtol=1e-3)
 
+
 if __name__ == "__main__":
   absltest.main()