jit(grad(f))(x) and grad(jit(f))(x) do not have same effect on the parameters

[rmoyard]

Hi Jax team! This behaviour has changed in 0.4.4, previously we would have exactly the same trace on the parameters by switching grad and jit but now we do not have access to the same information on the trace.

x = jax.numpy.array(0.1)

def h(x):
    print(x)
    return 3 * x

grad = jax.grad(jax.jit(h))(x)

grad = jax.jit(jax.grad(h))(x)

Traced<ShapedArray(float32[], weak_type=True)>with<DynamicJaxprTrace(level=3/0)>

Traced<ShapedArray(float32[], weak_type=True)>with<JVPTrace(level=3/0)> with
  primal = Traced<ShapedArray(float32[], weak_type=True)>with<DynamicJaxprTrace(level=1/0)>
  tangent = Traced<ShapedArray(float32[], weak_type=True)>with<JaxprTrace(level=2/0)> with
    pval = (ShapedArray(float32[], weak_type=True), None)
    recipe = LambdaBinding()

Is there still a way to access the JVPTrace on the first example? We are using isinstance(x, JVPTracer) to check trainability in a complicated workflow.

[jakevdp]

Can you say more about what you mean by "check trainability"? Tracer subclasses are an internal implementation detail, and isinstance checks against private APIs are not going to be a stable way to reason about a JAX program across JAX versions. But if you describe your intent, we may be able to recommend a public API that accomplishes the same thing.

[rmoyard]

Thanks for your answer! We use this JVPTracer check to get the parameters that were set as trainable by the user. Our library is PennyLane a quantum computing framework and we deal with quantum circuits, we use finite differences, parameter shift (quantum derivative) to get jacobian and we register JVPs for the quantum part. For example, let's say you have 4 parametrs and user set argnums=[0, 2] in jax.jacobian, we need this information to build the quantum jacobian and register the custom JVP. Without the arnums information we would have to always build the full jacobian and therefore have more executions on the quantum computer which is non-desirable.

Is there a way to access the information of what parameters were set trainable by the users (argnums)?

[froystig]

@mattjj, @patrick-kidger and I recently added symbolic zero support to jax.custom_jvp (#14570). Does that give you what you need?

Using this, if an array doesn't participate in differentiation, then it will enter as a symbolic zero, not take up any memory, and your custom JVP can handle things accordingly from there. Example:

from functools import partial
import jax

Zero = jax.custom_derivatives.SymbolicZero

@jax.custom_jvp
def f(x, y): return x * y

@partial(f.defjvp, symbolic_zeros=True)
def fjvp(primals, tangents):
    x, y = primals
    tx, ty = tangents
    if type(tx) is Zero and type(ty) is Zero:
        tz = None
    elif type(tx) is Zero:
        tz = x * ty
    elif type(ty) is Zero:
        tz = y * tx
    else:
        tz = x * ty + y * tx
    return f(x, y), tz

>>> jax.jacfwd(f, argnums=0)(2., 3.)
Array(3., dtype=float32, weak_type=True)
>>> jax.jacfwd(f, argnums=1)(2., 3.)
Array(2., dtype=float32, weak_type=True)
>>> jax.jacfwd(f, argnums=[0, 1])(2., 3.)
(Array(3., dtype=float32, weak_type=True), Array(2., dtype=float32, weak_type=True))

As @jakevdp wrote, traces and tracers are jax-internal implementation details.

[rmoyard]

@jakevdp @froystig I think this issue and the other one #14799 come from the switch to staged processing (initial style). The fact that symbolic zeros do not work when you jit first is probably also directly related to that.

Is there a way to get the argnums or symbolic zeros if you jit the function before taking the derivative?

[rmoyard]

From what I can gather the issue with symbolic zeros and jitting first arises in ad.py in process_custom_jvp_call, the tangents are processed and symbolic zeros replaces the zeros and then call the f_jvp,call_wrapped where the error arises. The SymbolicZeros object does not support self.pure in the trace.

  def process_custom_jvp_call(self, _, __, f_jvp, tracers, *, symbolic_zeros):
    primals_in, tangents_in = unzip2((t.primal, t.tangent) for t in tracers)
    primals_in = map(core.full_lower, primals_in)
    if not symbolic_zeros:
      tangents_in = map(instantiate_zeros, tangents_in)
      tangents_in = map(replace_float0s, primals_in, tangents_in)
    else:
      tangents_in = map(replace_internal_symbolic_zeros, tangents_in)
    outs = f_jvp.call_wrapped(*it.chain(primals_in, tangents_in))
    primals_out, tangents_out = split_list(outs, [len(outs) // 2])
    tangents_out = map(replace_rule_output_symbolic_zeros, tangents_out)
    tangents_out = map(recast_to_float0, primals_out, tangents_out)
    return map(partial(JVPTracer, self), primals_out, tangents_out)

[froystig]

Thanks! This is a bug and we should fix it. Filed #14833.

Symbolic zeros ought to work in all settings (absent any bugs). The best workaround in the meantime may be to avoid jitting, if that's feasible. See also jax.disable_jit if it helps.

[rmoyard]

Thanks @froystig, I am looking forward! Unfortunately we cannot disable jitting, we will pin to Jax 0.4.3 until we can get the argnums in a more official way hopefully symbolic zeros will cover all the cases.

[mattjj]

Fixed in 03e72e3!