# Buggy code
def _dot_general_shape_rule(lhs, rhs, dimension_numbers):
    (lhs_contracting, rhs_contracting), (lhs_batch, rhs_batch) = dimension_numbers
    if len(lhs_batch) != len(rhs_batch):
        msg = (
            "dot_general requires equal numbers of lhs_batch and rhs_batch "
            "dimensions, got lhs_batch {} and rhs_batch {}."
        )
        raise TypeError(msg.format(lhs_batch, rhs_batch))
    if not onp.all(onp.equal(lhs_batch, rhs_batch)):
        msg = (
            "dot_general requires same lhs and rhs batch dimension numbers, "
            "got {} and {}."
        )
        raise TypeError(msg.format(lhs_batch, rhs_batch))
    lhs_batch_shape = onp.take(lhs.shape, lhs_batch)
    rhs_batch_shape = onp.take(rhs.shape, rhs_batch)
    if not onp.all(onp.equal(lhs_batch_shape, rhs_batch_shape)):
        msg = (
            "dot_general requires lhs batch dimensions and rhs batch dimensions "
            "to have the same shape, got {} and {}."
        )
        raise TypeError(msg.format(lhs_batch_shape, rhs_batch_shape))
    if tuple(sorted(lhs_batch)) != tuple(range(len(lhs_batch))):
        msg = (
            "dot_general requires lhs batch dimensions to precede contracting "
            "and non-contracting dimensions, got lhs_batch {}."
        )
        raise TypeError(msg.format(lhs_batch))
    if tuple(sorted(rhs_batch)) != tuple(range(len(rhs_batch))):
        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)):
        msg = (
            "dot_general requires contracting dimensions to have the same "
            "shape, got {} and {}."
        )
        raise TypeError(msg.format(lhs_contracting_shape, rhs_contracting_shape))
    if lhs.ndim > len(lhs_batch) + len(lhs_contracting) + 1:
        msg = (
            "dot_general requires either one or zero non-batch non-contracting "
            "lhs dimension, got {}."
        )
        diff = lhs.ndim - len(lhs_batch) - len(lhs_contracting)
        raise TypeError(msg.format(diff))
    if rhs.ndim > len(rhs_batch) + len(rhs_contracting) + 1:
        msg = (
            "dot_general requires either one or zero non-batch non-contracting "
            "rhs dimension, got {}."
        )
        diff = rhs.ndim - len(rhs_batch) - len(rhs_contracting)
        raise TypeError(msg.format(diff))

    batch_shape = tuple(onp.take(lhs.shape, lhs_batch))
    lhs_contract_or_batch = tuple(lhs_contracting) + tuple(lhs_batch)
    lhs_tensored_shape = tuple(onp.delete(lhs.shape, lhs_contract_or_batch))
    rhs_contract_or_batch = tuple(rhs_contracting) + tuple(rhs_batch)
    rhs_tensored_shape = tuple(onp.delete(rhs.shape, rhs_contract_or_batch))
    return batch_shape + lhs_tensored_shape + rhs_tensored_shape


# Traceback (TypeError)
/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).

# Fixed code
def _dot_general_shape_rule(lhs, rhs, dimension_numbers):
    (lhs_contracting, rhs_contracting), (lhs_batch, rhs_batch) = dimension_numbers
    if len(lhs_batch) != len(rhs_batch):
        msg = (
            "dot_general requires equal numbers of lhs_batch and rhs_batch "
            "dimensions, got lhs_batch {} and rhs_batch {}."
        )
        raise TypeError(msg.format(lhs_batch, rhs_batch))
    if not onp.all(onp.equal(lhs_batch, rhs_batch)):
        msg = (
            "dot_general requires same lhs and rhs batch dimension numbers, "
            "got {} and {}."
        )
        raise TypeError(msg.format(lhs_batch, rhs_batch))
    lhs_batch_shape = onp.take(lhs.shape, lhs_batch)
    rhs_batch_shape = onp.take(rhs.shape, rhs_batch)
    if not onp.all(onp.equal(lhs_batch_shape, rhs_batch_shape)):
        msg = (
            "dot_general requires lhs batch dimensions and rhs batch dimensions "
            "to have the same shape, got {} and {}."
        )
        raise TypeError(msg.format(lhs_batch_shape, rhs_batch_shape))
    if tuple(sorted(lhs_batch)) != tuple(range(len(lhs_batch))):
        msg = (
            "dot_general requires lhs batch dimensions to precede contracting "
            "and non-contracting dimensions, got lhs_batch {}."
        )
        raise TypeError(msg.format(lhs_batch))
    if tuple(sorted(rhs_batch)) != tuple(range(len(rhs_batch))):
        msg = (
            "dot_general requires rhs batch dimensions to precede contracting "
            "and non-contracting dimensions, got rhs_batch {}."
        )
        raise TypeError(msg.format(rhs_batch))
    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)):
        msg = (
            "dot_general requires contracting dimensions to have the same "
            "shape, got {} and {}."
        )
        raise TypeError(msg.format(lhs_contracting_shape, rhs_contracting_shape))
    if lhs.ndim > len(lhs_batch) + len(lhs_contracting) + 1:
        msg = (
            "dot_general requires either one or zero non-batch non-contracting "
            "lhs dimension, got {}."
        )
        diff = lhs.ndim - len(lhs_batch) - len(lhs_contracting)
        raise TypeError(msg.format(diff))
    if rhs.ndim > len(rhs_batch) + len(rhs_contracting) + 1:
        msg = (
            "dot_general requires either one or zero non-batch non-contracting "
            "rhs dimension, got {}."
        )
        diff = rhs.ndim - len(rhs_batch) - len(rhs_contracting)
        raise TypeError(msg.format(diff))

    batch_shape = tuple(onp.take(lhs.shape, lhs_batch))
    lhs_contract_or_batch = tuple(lhs_contracting) + tuple(lhs_batch)
    lhs_tensored_shape = tuple(onp.delete(lhs.shape, lhs_contract_or_batch))
    rhs_contract_or_batch = tuple(rhs_contracting) + tuple(rhs_batch)
    rhs_tensored_shape = tuple(onp.delete(rhs.shape, rhs_contract_or_batch))
    return batch_shape + lhs_tensored_shape + rhs_tensored_shape