For jax.grad, why is it traced by a tracer of AbstractArray during jax.make_jaxpr?

[riven314]

I have a function (i.e. func) that I wanna transform by jax.grad. Meanwhile, I wanna inspect the type of the input during the tracing.

I noticed in normal call (e.g. jax.grad(func)(1.)) its traced by a tracer of ConcreteArray, but when its wrapped by jax.make_jaxpr (I wanna inspect the intermediate representation) its shown to be traced by a tracer of AbstractArray.
Here is an example:

# CODE AS ILLUSTRATION
def func(x):
    print(f"type of x: {x}")
    return jnp.sin(x)

grad_func = jax.grad(func)
grad_func(1.)
pr = jax.make_jaxpr(grad_func)(1.)

# OUTPUT THAT I RECEIVED
type of x: Traced<ConcreteArray(1.0, dtype=float32, weak_type=True)>with<JVPTrace(level=2/0)> with
  primal = Traced<ConcreteArray(1.0, dtype=float32, weak_type=True)>with<JaxprTrace(level=1/0)> with
    pval = (None, 1.0)
    recipe = *
  tangent = Traced<ShapedArray(float32[], weak_type=True)>with<JaxprTrace(level=1/0)> with
    pval = (ShapedArray(float32[], weak_type=True), *)
    recipe = LambdaBinding()
type of x: Traced<ShapedArray(float32[], weak_type=True)>with<JVPTrace(level=3/0)> with
  primal = Traced<ShapedArray(float32[], weak_type=True)>with<JaxprTrace(level=2/0)> with
    pval = (None, Traced<ShapedArray(float32[], weak_type=True)>with<DynamicJaxprTrace(level=1/0)>)
    recipe = *
  tangent = Traced<ShapedArray(float32[], weak_type=True)>with<JaxprTrace(level=2/0)> with
    pval = (ShapedArray(float32[], weak_type=True), *)
    recipe = LambdaBinding()

I wanna understand more about this discrepancy. Why in this case (i.e. jax.grad) different tracers (ConcreteArray v.s. AbstractArray) are used between normal call from the transformed function v.s. call from jax.make_jaxpr?

[jakevdp]

Two things to know: ShapedArray tracers basically contain references to the array's shape and dtype, while ConcreteArray tracers contain references to the array's shape, dtype, as well as its contents.

jax.grad uses concrete tracers in order to support value-based control flow during tracing. For example:

import jax

def f(x):
  if x < 0:
    return 0.0
  return x ** 2

print(jax.grad(f)(1.0))
# 2.0

If grad did not use concrete tracers in this case, the result of the condition x < 0 would not be knowable at trace-time, and we could not differentiate functions that use this kind of control flow.

By contrast, jaxprs are created entirely based on abstract properties of arrays (namely shape and dtype). This is for a number of reasons, including the fact that XLA compilation (i.e. jax.jit) is dispatched based on these abstract properties, and not the contents of arrays. For this case, a ShapedArray tracer contains all the relevant information, and it means that this kind of value-based dispatch is not allowed. For this reason, you'll find if you wrap the above function in make_jaxpr, the result will be a ConcretizationTypeError.

[riven314]

thanks so much for the address! it clears the confusion!

[riven314]

Hi @jakevdp ! Sorry for pinning you again even though I have marked the post as answered.
The reason is that I can't reconcile what I saw from documentation and what you answered (i.e. "autodiff does not trace to jaxprs as an intermediate step").

In JAX documentation, I saw it said:

• The linearize and vjp autodiff functions are built on jvp, but involve jaxprs as well. That’s because both involve staging out, or delaying, computation.
• To build the f_lin jaxpr from a JVP, we need to perform partial evaluation: we evaluate all the primal values as we trace, but stage the tangent computations into a jaxpr.
  (the above are from Part 4: linearize and vjp (and grad!)).

Also you made an excellent YT video introducing JAX. In the video you showed a diagram (at 8:41 of this video) illustrating the pipeline for function transformation. From the diagram, jax.grad involves JAX IR in the process.

So is it correct for me to say that jax.grad also involves jaxpr conversion, but instead it's doing a partial conversion which target part of the function that can "abstractably" traced by abstract tracer, unlike other function transformations that are doing a complete conversion to jaxpr?

[jakevdp]

There may be details of this that I'm not aware of. Perhaps someone else from the team can chime in regarding these docs.

[soraros]

@riven314 I'm not an expert on this and have only a read-through of the Autodidax not very long ago, but I'd try my best to answer your question.

1. Partial evaluation means a very specific thing here, as in this, since back prop (vjp and grad) requires it: only the primal part of evaluation can happen during tracing. Even if Jax doesn't use a tracing jit and python behaves more like a CAS, the tangent part of the computation can still only be delayed, as they have data dependency on the primal values that are just not available on the first run-through of the code(trying to avoid terminology like "tracing" or "staging" here).
2. It is also true that in order to make python control flow work, using abstract tracers along just won't cut it. Indeed, only part of the computation is abstractly traceable, but that is rather orthogonal to the kind of "partial" mentioned in (1).

Also, I think a good answer from the experts can get very long very easily, and in order to increase visibility, you might want to move your comment to a standalone question in the forum.

[riven314]

@soraros
thanks for clarifying the meaning of partial evaluation here. sure, I have opened up a new post for this discussion: #8817