Repository: google/jax

Issue:
Omnistaging breaks jax.scipy.sparse.linalg.cg in some settings
Not quite sure what is the issue but I boiled it down to the following.

```python
import jax
import jax.numpy as jnp
from jax.scipy.sparse.linalg import cg


def func(x):
    return x

x = jnp.array([-50., 200.])
val, func_jvp = jax.linearize(func, x)
cg(func_jvp, val)  # this fails on jax 0.2.0, but works with jax.config.disable_omnistaging()
```
I was able to reproduce this on a colab instance running jax 0.2.0.

<details><summary>The stack trace(s)</summary>
<p>

```
/usr/local/lib/python3.6/dist-packages/jax/lib/xla_bridge.py:130: UserWarning: No GPU/TPU found, falling back to CPU.
  warnings.warn('No GPU/TPU found, falling back to CPU.')
---------------------------------------------------------------------------
ConcretizationTypeError                   Traceback (most recent call last)
/usr/local/lib/python3.6/dist-packages/jax/api.py in fun(*tangents)
   1762       try:
-> 1763         core.lattice_join(primal_aval, tangent_aval)
   1764       except TypeError as e:

17 frames
/usr/local/lib/python3.6/dist-packages/jax/core.py in lattice_join(x, y)
    787   elif isinstance(x, type(y)):
--> 788     return y.join(x)
    789   elif isinstance(y, type(x)):

/usr/local/lib/python3.6/dist-packages/jax/core.py in join(self, other)
   1019   def join(self, other) -> UnshapedArray:
-> 1020     if self == other:
   1021       return self

/usr/local/lib/python3.6/dist-packages/jax/core.py in __bool__(self)
    506   def __nonzero__(self): return self.aval._nonzero(self)
--> 507   def __bool__(self): return self.aval._bool(self)
    508   def __int__(self): return self.aval._int(self)

/usr/local/lib/python3.6/dist-packages/jax/core.py in error(self, arg)
    863   def error(self, arg):
--> 864     raise_concretization_error(arg, fname_context)
    865   return error

/usr/local/lib/python3.6/dist-packages/jax/core.py in raise_concretization_error(val, context)
    852           f"Encountered tracer value: {val}")
--> 853   raise ConcretizationTypeError(msg)
    854 

ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected.

The problem arose with the `bool` function. 

While tracing the function f at /usr/local/lib/python3.6/dist-packages/jax/lax/lax_control_flow.py:2085, this value became a tracer due to JAX operations on these lines:

  operation m:bool[2] = eq k:float32[2] l:float32[2]
    from line <ipython-input-1-fdc0af3e82fa>:12 (<module>)

You can use transformation parameters such as `static_argnums` for `jit` to avoid tracing particular arguments of transformed functions.

See https://jax.readthedocs.io/en/latest/faq.html#abstract-tracer-value-encountered-where-concrete-value-is-expected-error for more information.

Encountered tracer value: Traced<ShapedArray(bool[])>with<DynamicJaxprTrace(level=1/0)>

The above exception was the direct cause of the following exception:

ValueError                                Traceback (most recent call last)
<ipython-input-1-fdc0af3e82fa> in <module>()
     10 
     11 val, func_jvp = jax.linearize(func, x)
---> 12 print(cg(func_jvp, val))

/usr/local/lib/python3.6/dist-packages/jax/scipy/sparse/linalg.py in cg(A, b, x0, tol, atol, maxiter, M)
    172   symmetric = all(map(real_valued, tree_leaves(b)))
    173   x = lax.custom_linear_solve(
--> 174       A, b, solve=cg_solve, transpose_solve=cg_solve, symmetric=symmetric)
    175   info = None  # TODO(shoyer): return the real iteration count here
    176   return x, info

/usr/local/lib/python3.6/dist-packages/jax/lax/lax_control_flow.py in custom_linear_solve(matvec, b, solve, transpose_solve, symmetric)
   2094 
   2095   solve_jaxpr, solve_consts, out_tree = _initial_style_jaxpr(
-> 2096       _shape_checked(partial(solve, matvec), "solve"), in_args_tree, b_avals)
   2097   _check_tree("solve", "b", out_tree, tree)
   2098 

/usr/local/lib/python3.6/dist-packages/jax/lax/lax_control_flow.py in _initial_style_jaxpr(fun, in_tree, in_avals)
     70 @cache()
     71 def _initial_style_jaxpr(fun: Callable, in_tree, in_avals):
---> 72   jaxpr, out_avals, consts, out_tree = _initial_style_open_jaxpr(fun, in_tree, in_avals)
     73   closed_jaxpr = core.ClosedJaxpr(pe.convert_constvars_jaxpr(jaxpr), ())
     74   return closed_jaxpr, consts, out_tree

/usr/local/lib/python3.6/dist-packages/jax/lax/lax_control_flow.py in _initial_style_open_jaxpr(fun, in_tree, in_avals)
     65 def _initial_style_open_jaxpr(fun: Callable, in_tree, in_avals):
     66   wrapped_fun, out_tree = flatten_fun_nokwargs(lu.wrap_init(fun), in_tree)
---> 67   jaxpr, out_avals, consts = pe.trace_to_jaxpr_dynamic(wrapped_fun, in_avals)
     68   return jaxpr, out_avals, consts, out_tree()
     69 

/usr/local/lib/python3.6/dist-packages/jax/interpreters/partial_eval.py in trace_to_jaxpr_dynamic(fun, in_avals)
    992     main.source_info = fun_sourceinfo(fun.f)  # type: ignore
    993     main.jaxpr_stack = ()  # type: ignore
--> 994     jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(fun, main, in_avals)
    995     del main
    996   return jaxpr, out_avals, consts

/usr/local/lib/python3.6/dist-packages/jax/interpreters/partial_eval.py in trace_to_subjaxpr_dynamic(fun, main, in_avals)
   1002     trace = DynamicJaxprTrace(main, core.cur_sublevel())
   1003     in_tracers = map(trace.new_arg, in_avals)
-> 1004     ans = fun.call_wrapped(*in_tracers)
   1005     out_tracers = map(trace.full_raise, ans)
   1006   jaxpr, out_avals, consts = frame.to_jaxpr(in_tracers, out_tracers)

/usr/local/lib/python3.6/dist-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    149 
    150     try:
--> 151       ans = self.f(*args, **dict(self.params, **kwargs))
    152     except:
    153       # Some transformations yield from inside context managers, so we have to

/usr/local/lib/python3.6/dist-packages/jax/lax/lax_control_flow.py in f(x)
   2084   def _shape_checked(fun, name):
   2085     def f(x):
-> 2086       y = fun(x)
   2087       _check_shapes(name, "b", y, b_flat)
   2088       return y

/usr/local/lib/python3.6/dist-packages/jax/scipy/sparse/linalg.py in _cg_solve(A, b, x0, maxiter, tol, atol, M)
     77     return x_, r_, gamma_, p_, k + 1
     78 
---> 79   r0 = _sub(b, A(x0))
     80   p0 = z0 = M(r0)
     81   gamma0 = _vdot_tree(r0, z0)

/usr/local/lib/python3.6/dist-packages/jax/api.py in _lift_linearized(jaxpr, primal_avals, consts, io_tree, out_pvals, *py_args)
   1770                      out_pvals, tangents_out))
   1771 
-> 1772   return apply_flat_fun(fun, io_tree, *py_args)
   1773 
   1774 def _check_inexact_input_vjp(x):

/usr/local/lib/python3.6/dist-packages/jax/api_util.py in apply_flat_fun(fun, io_tree, *py_args)
     50   if in_tree != in_tree_expected:
     51       raise TypeError("Expected {}, got {}".format(in_tree_expected, in_tree))
---> 52   ans = fun(*args)
     53   return tree_unflatten(out_tree, ans)
     54 

/usr/local/lib/python3.6/dist-packages/jax/api.py in fun(*tangents)
   1765         msg = ("linearized function called on tangent values inconsistent with "
   1766                "the original primal values.")
-> 1767         raise ValueError(msg) from e
   1768     tangents_out = eval_jaxpr(jaxpr, consts, *tangents)
   1769     return tuple(map(lambda out_pv, tan_out: out_pv.merge_with_known(tan_out),

ValueError: linearized function called on tangent values inconsistent with the original primal values.
```
</p>
</details>



Discussion:

Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -1762,11 +1762,12 @@ def fun(*tangents):
     tangent_avals = list(map(core.get_aval, tangents))
     for primal_aval, tangent_aval in zip(primal_avals, tangent_avals):
       try:
-        core.lattice_join(primal_aval, tangent_aval)
+        core.lattice_join(primal_aval.at_least_vspace(), tangent_aval)
       except TypeError as e:
         msg = ("linearized function called on tangent values inconsistent with "
-               "the original primal values.")
-        raise ValueError(msg) from e
+               "the original primal values: "
+               f"got {tangent_aval} for primal aval {primal_aval}")
+        raise ValueError(msg)
     tangents_out = eval_jaxpr(jaxpr, consts, *tangents)
     return tuple(map(lambda out_pv, tan_out: out_pv.merge_with_known(tan_out),
                      out_pvals, tangents_out))
diff --git a/jax/core.py b/jax/core.py
--- a/jax/core.py
+++ b/jax/core.py
@@ -1025,9 +1025,12 @@ def __init__(self, val, weak_type=False):
     assert self.dtype != np.dtype('O'), val
 
   def __eq__(self, other):
-    return (type(self) is type(other) and self.dtype == other.dtype
-            and self.shape == other.shape and self.weak_type == other.weak_type
-            and np.all(self.val == other.val))
+    if (type(self) is type(other) and self.dtype == other.dtype
+        and self.shape == other.shape and self.weak_type == other.weak_type):
+      with eval_context():  # in case self.val is a DeviceArray
+        return (self.val == other.val).all()
+    else:
+      return False
 
   def __hash__(self):
     return id(self.val)


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -1133,16 +1133,16 @@ def __init__(self, shape, dtype):
   def test_issue_871(self):
     T = jnp.array([[1., 2.], [3., 4.], [5., 6.]])
     x = jnp.array([1, 2, 3])
+    msg = ("linearized function called on tangent values inconsistent with "
+           "the original primal values")
 
     y, f_jvp = api.linearize(jnp.sum, x)
-    jtu.check_raises(lambda: f_jvp(T), ValueError,
-                     ("linearized function called on tangent values "
-                      "inconsistent with the original primal values."))
+    with self.assertRaisesRegex(ValueError, msg):
+      f_jvp(T)
 
     y, f_jvp = api.linearize(api.jit(jnp.sum), x)
-    jtu.check_raises(lambda: f_jvp(T), ValueError,
-                     ("linearized function called on tangent values "
-                      "inconsistent with the original primal values."))
+    with self.assertRaisesRegex(ValueError, msg):
+      f_jvp(T)
 
   def test_partial_eval_lower(self):
     # this is a simplified model of a bug that arose when we first used @jit in
@@ -1966,6 +1966,38 @@ def device_put_and_count(*args, **kwargs):
       xla.device_put = orig_device_put
     self.assertEqual(count, 0)
 
+  def test_join_concrete_arrays_with_omnistaging(self):
+    # https://github.com/google/jax/issues/4622
+    if not config.omnistaging_enabled:
+      raise unittest.SkipTest("test is omnistaging-specific")
+
+    x = jnp.array([1., 2., 3.])
+    y = jnp.array([1., 2., 4.])
+
+    @jit
+    def f():
+      core.lattice_join(core.ConcreteArray(x), core.ConcreteArray(y))
+
+    f()  # doesn't crash
+
+  def test_linearize_aval_error(self):
+    # https://github.com/google/jax/issues/4622
+    f = lambda x: x
+
+    # these should not error
+    _, f_jvp = api.linearize(f, 1.)
+    f_jvp(1.)
+    _, f_jvp = api.linearize(f, np.ones(2, np.int32))
+    f_jvp(np.zeros(2, float0))
+
+    # these should error
+    _, f_jvp = api.linearize(f, 1.)
+    with self.assertRaisesRegex(ValueError, "tangent values inconsistent"):
+      f_jvp(1)
+    _, f_jvp = api.linearize(f, np.ones(2, np.int32))
+    with self.assertRaisesRegex(ValueError, "tangent values inconsistent"):
+      f_jvp(np.ones(2, np.int32))
+
 
 class RematTest(jtu.JaxTestCase):