Repository: google/jax
Issue: Hessian calculation finds an UnshapedArray when jitted.

The code below throws an error that it finds an unshaped array only when f is jitted and the line `dR = dR - np.where(np.abs(dR) < 0.5, zero, 0.5 * np.sign(dR))` is present. 

```
R = npr.rand(100, 2)
def dist_sq(R):
  dR = R[:, np.newaxis, :] - R[np.newaxis, :, :]
  zero = np.zeros_like(dR)
  dR = dR - np.where(np.abs(dR) < 0.5, zero, 0.5 * np.sign(dR))
  return np.sum(dR ** 2, axis=2)

@jit
def f(R):
  dr = dist_sq(R)
  return np.sum(R ** 2)

H = hessian(f)(R)
```

Here is the full trace.

```
google3/third_party/py/jax/api.py in jacfun(*args, **kwargs)
    183     f_partial, dyn_args = argnums_partial(f, argnums, args)
    184     pushfwd = partial(jvp, f_partial, dyn_args)
--> 185     y, jac = vmap(pushfwd, out_axes=(None, -1))(_std_basis(dyn_args))
    186     example_args = dyn_args[0] if isinstance(argnums, int) else dyn_args
    187     return tree_map(partial(_unravel_array_into_pytree, example_args, -1), jac)

google3/third_party/py/jax/api.py in batched_fun(*args, **kwargs)
    273     in_flat, in_trees = unzip2(map(pytree_to_jaxtupletree, args))
    274     jaxtree_fun, out_tree = pytree_fun_to_jaxtupletree_fun(f, in_trees)
--> 275     out_flat = batching.batch(jaxtree_fun, in_flat, in_axes_, out_axes)
    276     return build_tree(out_tree(), out_flat)
    277 

google3/third_party/py/jax/interpreters/batching.py 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:

google3/third_party/py/jax/linear_util.py in call_wrapped(self, *args)
     84 
     85     del gen
---> 86     ans = self.f(*args, **self.kwargs)
     87     del args
     88     while stack:

google3/third_party/py/jax/api.py in jvp(fun, primals, tangents)
    350   ps_flat, ts_flat, in_trees = unzip3(map(trim_arg, primals, tangents))
    351   jaxtree_fun, out_tree = pytree_fun_to_jaxtupletree_fun(fun, in_trees)
--> 352   out_primal, out_tangent = ad.jvp(jaxtree_fun).call_wrapped(ps_flat, ts_flat)
    353   return (build_tree(out_tree(), out_primal), build_tree(out_tree(), out_tangent))
    354 

google3/third_party/py/jax/linear_util.py in call_wrapped(self, *args)
     84 
     85     del gen
---> 86     ans = self.f(*args, **self.kwargs)
     87     del args
     88     while stack:

google3/third_party/py/jax/api.py in jacfun(*args, **kwargs)
    195     f = lu.wrap_init(fun, kwargs)
    196     f_partial, dyn_args = argnums_partial(f, argnums, args)
--> 197     y, pullback = vjp(f_partial, *dyn_args)
    198     jac = vmap(pullback)(_std_basis(y))
    199     jac = jac[0] if isinstance(argnums, int) else jac

google3/third_party/py/jax/api.py in vjp(fun, *primals)
    378   check_args(primals_flat)
    379   jaxtree_fun, out_tree = pytree_fun_to_jaxtupletree_fun(fun, in_trees)
--> 380   out_primal, out_vjp = ad.vjp(jaxtree_fun, primals_flat)
    381   out_tree = out_tree()
    382   out_primal_py = build_tree(out_tree, out_primal)

google3/third_party/py/jax/interpreters/ad.py 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)

google3/third_party/py/jax/interpreters/ad.py 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

google3/third_party/py/jax/interpreters/partial_eval.py 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

google3/third_party/py/jax/linear_util.py in call_wrapped(self, *args)
     84 
     85     del gen
---> 86     ans = self.f(*args, **self.kwargs)
     87     del args
     88     while stack:

google3/third_party/py/jax/api.py in f_jitted(*args, **kwargs)
     89     check_args(jaxtupletree_args)
     90     jaxtree_fun, out_tree = pytree_fun_to_jaxtupletree_fun(f, in_trees)
---> 91     jaxtupletree_out = xla.xla_call(jaxtree_fun, *jaxtupletree_args)
     92     return build_tree(out_tree(), jaxtupletree_out)
     93 

google3/third_party/py/jax/core.py in call_bind(primitive, f, *args, **kwargs)
    534   else:
    535     tracers = map(top_trace.full_raise, args)
--> 536     ans = full_lower(top_trace.process_call(primitive, f, tracers, kwargs))
    537   return apply_todos(env_trace_todo(), ans)
    538 

google3/third_party/py/jax/interpreters/ad.py in process_call(self, call_primitive, f, tracers, params)
    191       result = call_primitive.bind(f, pack(primals), nonzero_tangents, **new_params)
    192     else:
--> 193       result = call_primitive.bind(f, pack(primals), nonzero_tangents, **params)
    194     primal_out, tangent_out = build_tree(out_tree_def(), result)
    195     return JVPTracer(self, primal_out, tangent_out)

google3/third_party/py/jax/core.py in call_bind(primitive, f, *args, **kwargs)
    534   else:
    535     tracers = map(top_trace.full_raise, args)
--> 536     ans = full_lower(top_trace.process_call(primitive, f, tracers, kwargs))
    537   return apply_todos(env_trace_todo(), ans)
    538 

google3/third_party/py/jax/interpreters/partial_eval.py in process_call(self, call_primitive, f, tracers, params)
     79     in_pvs, in_consts = unzip2([t.pval for t in tracers])
     80     fun, aux = partial_eval(f, self, in_pvs)
---> 81     out_pv_const, consts = call_primitive.bind(fun, *in_consts, **params)
     82     out_pv, jaxpr, env = aux()
     83     const_tracers = map(self.new_instantiated_const, consts)

google3/third_party/py/jax/core.py in call_bind(primitive, f, *args, **kwargs)
    534   else:
    535     tracers = map(top_trace.full_raise, args)
--> 536     ans = full_lower(top_trace.process_call(primitive, f, tracers, kwargs))
    537   return apply_todos(env_trace_todo(), ans)
    538 

google3/third_party/py/jax/interpreters/ad.py in process_call(self, call_primitive, f, tracers, params)
    191       result = call_primitive.bind(f, pack(primals), nonzero_tangents, **new_params)
    192     else:
--> 193       result = call_primitive.bind(f, pack(primals), nonzero_tangents, **params)
    194     primal_out, tangent_out = build_tree(out_tree_def(), result)
    195     return JVPTracer(self, primal_out, tangent_out)

google3/third_party/py/jax/core.py in call_bind(primitive, f, *args, **kwargs)
    534   else:
    535     tracers = map(top_trace.full_raise, args)
--> 536     ans = full_lower(top_trace.process_call(primitive, f, tracers, kwargs))
    537   return apply_todos(env_trace_todo(), ans)
    538 

google3/third_party/py/jax/interpreters/batching.py in process_call(self, call_primitive, f, tracers, params)
    129     else:
    130       f, dim_out = batch_subtrace(f, self.master, dims)
--> 131       val_out = call_primitive.bind(f, *vals, **params)
    132       return BatchTracer(self, val_out, dim_out())
    133 

google3/third_party/py/jax/core.py in call_bind(primitive, f, *args, **kwargs)
    531   if top_trace is None:
    532     with new_sublevel():
--> 533       ans = primitive.impl(f, *args, **kwargs)
    534   else:
    535     tracers = map(top_trace.full_raise, args)

google3/third_party/py/jax/interpreters/xla.py in xla_call_impl(fun, *args)
    420   fun, out_tree = flatten_fun(fun, in_trees)
    421 
--> 422   compiled_fun = xla_callable(fun, *map(abstractify, flat_args))
    423   flat_ans = compiled_fun(*flat_args)
    424 

google3/third_party/py/jax/linear_util.py in memoized_fun(f, *args)
    144       if len(cache) > max_size:
    145         cache.popitem(last=False)
--> 146       ans = call(f, *args)
    147       cache[key] = (ans, f)
    148     return ans

google3/third_party/py/jax/interpreters/xla.py in xla_callable(fun, *abstract_args)
    432   pvals = [PartialVal((aval, core.unit)) for aval in abstract_args]
    433   with core.new_master(JaxprTrace, True) as master:
--> 434     jaxpr, (pval, consts, env) = trace_to_subjaxpr(fun, master).call_wrapped(pvals)
    435     assert not env  # no subtraces here (though cond might eventually need them)
    436     compiled, result_shape = compile_jaxpr(jaxpr, consts, *abstract_args)

google3/third_party/py/jax/linear_util.py in call_wrapped(self, *args)
     84 
     85     del gen
---> 86     ans = self.f(*args, **self.kwargs)
     87     del args
     88     while stack:

<ipython-input-72-ce1d436081cc> in f(R)
      8 @jit
      9 def f(R):
---> 10   dr = dist_sq(R)
     11   return np.sum(R ** 2)
     12 

<ipython-input-72-ce1d436081cc> in dist_sq(R)
      4   zero = np.zeros_like(dR)
      5   dR = dR - np.where(np.abs(dR) < 0.5, zero, 0.5 * np.sign(dR))
----> 6   return np.sum(dR ** 2, axis=2)
      7 
      8 @jit

google3/third_party/py/jax/core.py in __pow__(self, other)
    231   def __mod__(self, other): return self.aval._mod(self, other)
    232   def __rmod__(self, other): return self.aval._rmod(self, other)
--> 233   def __pow__(self, other): return self.aval._pow(self, other)
    234   def __rpow__(self, other): return self.aval._rpow(self, other)
    235   def __matmul__(self, other): return self.aval._matmul(self, other)

google3/third_party/py/jax/numpy/lax_numpy.py in <lambda>(x, y)
    243 def _one_to_one_binop(numpy_fn, lax_fn, promote_like=False):
    244   if promote_like:
--> 245     fn = lambda x, y: lax_fn(*_promote_args_like(numpy_fn, x, y))
    246   else:
    247     fn = lambda x, y: lax_fn(*_promote_args(numpy_fn.__name__, x, y))

google3/third_party/py/jax/lax.py in pow(x, y)
    103 def conj(x): return conj_p.bind(x, input_dtype=_dtype(x))
    104 def abs(x): return abs_p.bind(x)
--> 105 def pow(x, y): return pow_p.bind(x, y)
    106 
    107 def bitwise_not(x): return not_p.bind(x)

google3/third_party/py/jax/core.py in bind(self, *args, **kwargs)
     76 
     77     tracers = map(top_trace.full_raise, args)
---> 78     out_tracer = top_trace.process_primitive(self, tracers, kwargs)
     79     return full_lower(out_tracer)
     80 

google3/third_party/py/jax/interpreters/ad.py 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 

google3/third_party/py/jax/interpreters/ad.py in standard_jvp(jvprules, primitive, primals, tangents, **params)
    299 
    300 def standard_jvp(jvprules, primitive, primals, tangents, **params):
--> 301   val_out = primitive.bind(*primals, **params)
    302   tangents_out = (rule(t, *primals, **params) for rule, t in zip(jvprules, tangents)
    303                   if rule is not None and t is not zero)

google3/third_party/py/jax/core.py in bind(self, *args, **kwargs)
     76 
     77     tracers = map(top_trace.full_raise, args)
---> 78     out_tracer = top_trace.process_primitive(self, tracers, kwargs)
     79     return full_lower(out_tracer)
     80 

google3/third_party/py/jax/interpreters/ad.py 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 

google3/third_party/py/jax/interpreters/ad.py in standard_jvp(jvprules, primitive, primals, tangents, **params)
    302   tangents_out = (rule(t, *primals, **params) for rule, t in zip(jvprules, tangents)
    303                   if rule is not None and t is not zero)
--> 304   return val_out, reduce(add_tangents, tangents_out, zero)
    305 
    306 

google3/third_party/py/jax/interpreters/ad.py in <genexpr>((rule, t))
    301   val_out = primitive.bind(*primals, **params)
    302   tangents_out = (rule(t, *primals, **params) for rule, t in zip(jvprules, tangents)
--> 303                   if rule is not None and t is not zero)
    304   return val_out, reduce(add_tangents, tangents_out, zero)
    305 

google3/third_party/py/jax/lax.py in <lambda>(g, x, y)
    959 pow_p = standard_binop([_float | _complex, _float | _complex], 'pow')
    960 ad.defjvp(pow_p,
--> 961           lambda g, x, y: mul(_brcast(g, y), mul(y, pow(x, select(
    962               eq(y, _zeros(y)), _ones(y), sub(y, _ones(y)))))),
    963           lambda g, x, y: mul(_brcast(g, x),

google3/third_party/py/jax/lax.py in _brcast(x, *others)
    836   # We don't need full numpy broadcasting, but otherwise the logic is the same
    837   # so we reuse the broadcast_shapes function after filtering out scalars.
--> 838   shapes = tuple(filter(None, map(onp.shape, (x,) + others)))
    839   shape = shapes and broadcast_shapes(*shapes)
    840   if onp.shape(x) != shape:

google3/third_party/py/numpy/core/fromnumeric.py in shape(a)
   1602     """
   1603     try:
-> 1604         result = a.shape
   1605     except AttributeError:
   1606         result = asarray(a).shape

google3/third_party/py/jax/core.py in __getattr__(self, name)
    262 
    263     try:
--> 264       attr = getattr(self.aval, name)
    265     except KeyError:
    266       raise AttributeError(

google3/third_party/py/jax/interpreters/batching.py in aval(self)
     82   @property
     83   def aval(self):
---> 84     batched_aval = get_aval(self.val)
     85     return remove_batch_dim_from_aval(self.batch_dim, batched_aval)
     86 

google3/third_party/py/jax/interpreters/batching.py in get_aval(x)
    152 def get_aval(x):
    153   if isinstance(x, Tracer):
--> 154     return raise_to_shaped(x.aval)
    155   else:
    156     return shaped_aval(x)

google3/third_party/py/jax/interpreters/batching.py in raise_to_shaped(aval)
    168     return ShapedArray(aval.shape, aval.dtype)
    169   else:
--> 170     raise TypeError(type(aval))
    171 
    172 def remove_batch_dim_from_aval(bdim, aval):

TypeError: <class 'jax.abstract_arrays.UnshapedArray'>
```

Changed files:
tests/batching_test.py, jax/interpreters/batching.py

diff --git a/jax/interpreters/batching.py b/jax/interpreters/batching.py
index fc8610d7e..d2eb068df 100644
--- a/jax/interpreters/batching.py
+++ b/jax/interpreters/batching.py
@@ -241,11 +241,13 @@ def reducer_batcher(prim, batched_args, batch_dims, axes, **params):
 
 def add_batched(batched_args, batch_dims):
   bdx, bdy = batch_dims
+  xs, ys = batched_args
   if bdx == bdy:
-    xs, ys = batched_args
     return add_jaxvals_p.bind(xs, ys), bdx
   else:
-    xs, ys = map(bdim_at_front, batched_args, batch_dims)
+    sz = (dimsize(bdx, xs) | dimsize(bdy, ys)).pop()
+    move_bdim = partial(bdim_at_front, broadcast_size=sz, force_broadcast=True)
+    xs, ys = map(move_bdim, batched_args, batch_dims)
     return add_jaxvals_p.bind(xs, ys), 0
 primitive_batchers[add_jaxvals_p] = add_batched
 
diff --git a/tests/batching_test.py b/tests/batching_test.py
index c573b57d0..e65071b7a 100644
--- a/tests/batching_test.py
+++ b/tests/batching_test.py
@@ -810,6 +810,23 @@ class BatchingTest(jtu.JaxTestCase):
     expected = onp.stack([grad(f)(scale) for scale in scales])
     self.assertAllClose(ans, expected, check_dtypes=False)
 
+  def testIssue387(self):
+    # https://github.com/google/jax/issues/387
+    R = onp.random.RandomState(0).rand(100, 2)
+
+    def dist_sq(R):
+      dR = R[:, np.newaxis, :] - R[np.newaxis, :, :]
+      zero = np.zeros_like(dR)
+      dR = dR - np.where(np.abs(dR) < 0.5, zero, 0.5 * np.sign(dR))
+      return np.sum(dR ** 2, axis=2)
+
+    @jit
+    def f(R):
+      dr = dist_sq(R)
+      return np.sum(R ** 2)
+
+    H = hessian(f)(R)  # don't crash on UnshapedArray
+
 
 if __name__ == '__main__':
   absltest.main()