# Buggy code
def _lift_linearized(jaxpr, primal_avals, consts, io_tree, out_pvals, *py_args):
    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)
            except TypeError as e:
                msg = (
                    "linearized function called on tangent values inconsistent with "
                    "the original primal values."
                )
                raise ValueError(msg) from e
        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,
            )
        )

    return apply_flat_fun(fun, io_tree, *py_args)


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

# Fixed code
def _lift_linearized(jaxpr, primal_avals, consts, io_tree, out_pvals, *py_args):
    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.at_least_vspace(), tangent_aval)
            except TypeError as e:
                msg = (
                    "linearized function called on tangent values inconsistent with "
                    "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,
            )
        )

    return apply_flat_fun(fun, io_tree, *py_args)