Is there a way to get JIT to treat constant inner function outputs as compile time constants?

[leakec]

One of the common patterns I use in JAX is the following:

f = lambda x,xi: np.dot(H(x),xi)
df = grad(f)
L = jit(lambda xi: df(x0,xi))

Essentially, f(x,xi) is some potentially complex function in x whose gradient with respect to x I don't want to compute manually (I have used the simple H(x) here, but in general it is not so simple). The grad transform would essentially give me something like df(x,xi) = np.dot(dH(x),xi) where dH(x) is the gradient of H(x). When calculating L I want to JIT compile df for a static-known value of x (called x0 in this example), but I would like the JIT compiler to compute and cache the output of dH(x0) (and in general cache and store the outputs of any functions that only depend on x). This would prevent L from having to re-compute dH(x) each time it is called, and greatly improve the run-time of L.

Is it possible to accomplish this with the tools JAX offers?

[YouJiacheng]

maybe you should check whether jit is already do the pre-computation for you?

lowered = L.lower(example_xi)
print(lowered.compile().compiler_ir()[0].to_string())

If not, you can define custom vjp for f, and explicit use pre-computed vjp for H.

[leakec]

In the code I am using this in H(x) is a JAX primitive that I have defined complete with its own defined XLA transform. In the XLA transform I added a print statement so I could check whether this function was being called every time the jitted function was being called and it is.

In generally, f(x,xi) won't be as simple as np.dot(H(x),xi), it could be arbitrarily complex, e.g., f(x,xi) = np.dot(H(x),xi) + x**2*np.dot(H(x),xi) + .... I could do a custom vjp for f(x,xi) as well, but L uses df. I would like to somehow precompute the all of the H(x)-like functions in df, even though in general I won't have df explicitly written out.

[leakec]

I ran the code you suggested and obtained the following result:

HloModule jit__lambda_.9

%fused_computation (param_0: f64[3,6], param_1.2: s64[]) -> f64[3,6] {
  %param_0 = f64[3,6]{1,0} parameter(0)
  %param_1.2 = s64[] parameter(1)
  %convert.0 = f64[] convert(s64[] %param_1.2), metadata={op_name="jit(<lambda>)/jit(main)/convert_element_type[new_dtype=float64 weak_type=False]" source_file="test.py" source_line=14}
  %broadcast.0 = f64[3,6]{1,0} broadcast(f64[] %convert.0), dimensions={}, metadata={op_name="jit(<lambda>)/jit(main)/mul" source_file="test.py" source_line=14}
  ROOT %multiply.0 = f64[3,6]{1,0} multiply(f64[3,6]{1,0} %param_0, f64[3,6]{1,0} %broadcast.0), metadata={op_name="jit(<lambda>)/jit(main)/mul" source_file="test.py" source_line=14}
}

ENTRY %main.17 (Arg_0.1: s64[]) -> f64[3,6] {
  %constant.9 = s64[] constant(0)
  %constant.2 = f64[3]{0} constant({0, 0.5, 1})
  %constant.10 = pred[] constant(false)
  %constant.11 = s64[] constant(3)
  %constant.12 = s64[] constant(6)
  %custom-call.1 = f64[3,6]{1,0} custom-call(s64[] %constant.9, f64[3]{0} %constant.2, s64[] %constant.9, pred[] %constant.10, s64[] %constant.11, /*index=5*/s64[] %constant.12), custom_call_target="BasisFunc0", metadata={source_file="-" source_line=9}
  %Arg_0.1 = s64[] parameter(0)
  ROOT %fusion = f64[3,6]{1,0} fusion(f64[3,6]{1,0} %custom-call.1, s64[] %Arg_0.1), kind=kLoop, calls=%fused_computation, metadata={op_name="jit(<lambda>)/jit(main)/mul" source_file="test.py" source_line=14}
}

At the end you can see the call to %custom-call.1 = ... custom_call_target="BasisFunc0 ... which is the H(x) function.

[YouJiacheng]

I suspect that a explicit partial evaluation is needed, since the precompute maybe very expensive and compiler should not be such aggressive.
For custom vjp, I mean you can define custom vjp for the composite function using vjp of "sub function".
Namely for every function f = lambda x,xi： fun(x, H(x), xi), where fun can be arbitrarily complex. I think you can define vjp for f using vjp for fun and vjp for H, and precompute the latter.

[YouJiacheng]

https://github.com/google/jax/blob/main/jax/interpreters/partial_eval.py
I think this tool can help you gracefully achieve your goal.
jax use it to implement linearize as follow:

You can declare x as known value then use pe.trace_to_jaxpr.
WDYT?

[mariogeiger]

@leakec would the following example be a useful and correct usage of pe?

@jax.jit
def f(x):
    y = compute_y(x)
    
    @pe(x, y, constant_arg_nums=[1])
    def g(x, y):
        # ...

    return g(x) + g(x**2) + g(x**3)

[leakec]

I don't think the way pe is designed will allow it to work on a function, defined inside a function that is being jitted. The reason is that jit is going to pass Tracer arrays during the compilation, but the arguments going into pe should be values of the same type/shape as those that will be passed into the real function, and in the case of constant_arg_nums should be the values you want to use exactly.

Also, how much speed gain you will see by using pe on g(x,y) depends on how y is used inside the function. However, I'm guessing compute_y is the expensive function here, and it will always be called once, regardless of whether pe is used or not.

[leakec]

I just tried a function like the one you suggested and the issue I mentioned in the previous response happened: pe does not work if it is on a function that is defined as an inner-function of a function with the jit decorator on it.

[mariogeiger]

Thanks for the clarifications

[leakec]

You're welcome

[PhilipVinc]

jax.ensure_compile_time_aval does exactly what you are asking

[YouJiacheng]

I think it is hard to find where to insert this context manager.
I believe everything inside it should be evaluated during compile time.
But in this case, things are interleaving, some can be evaluate during compile time while others not.

[YouJiacheng]

notably, the gradient w.r.t. x varies according to xi. Thus the "constant inner function" is not equivalent to a constant.

[YouJiacheng]

If you don't need the gradient w.r.t. x, you can just extract the evaluation of H(x0) out of f or follow @PhilipVinc 's suggestion.