Why aren't jitted function a JAX object?

[AdrienCorenflos]

Hi,
There are a number of cases when I have wanted to pass a jitted function as an argument to another one, but this is not possible (yet), contrarily to say tensorflow. Des this have something to do with possible closure problems/gradient tracing?
Would it possibly considered as a future improvement or is it a no go?

To fix ideas, I'm talking of doing something like this:

from jax import jit

@jit
def f(x):
    return x

@jit
def g(h, y):
    return h(y)

g(f, 5.)

[mattjj]

Does this do what you want?

from functools import partial
from jax import jit

@jit
def f(x):
    return x

@partial(jit, static_argnums=0)
def g(h, y):
    return h(y)

g(f, 5.)

[AdrienCorenflos]

Yeah so this is what I usually end up doing, but my problem is that I fail to understand why it is necessary. if h is compiled already, then why do we need to recompile everything that happens around h in g every time you pass in a new version.

Is it because you want to control the overall graph structure under the jit and can't allow for stitching-in predefined graphs?

[mattjj]

Ah, if I understand correctly, you're thinking of something like:

@jit
def f1(x):
    ...

@jit
def f2(x):
    ...

@partial(jit, static_argnums=?)
def g(h, y):
    return h(y)

g(f1, 5.)
g(f2, 3.)  # can we reuse parts of `g` from the previous call?

Is that right?

This may be a bit academic because XLA has no way to reuse pieces like this (i.e. it's not just a JAX limitation), though one can imagine compilers that could. The trouble is that those compilers would have much less information available when optimizing: for example, if g is not specialized on its function-valued argument, then the compiler doesn't know how much memory the function-valued argument might need to execute, or what input/output memory layouts are best, or whether parts of the function should be fused into consumers or rematerialized. In other words, it'd foil many of the optimizations that XLA is all about!

Are you concerned about a specific metric (e.g. compile times), or is this mainly an aesthetic question?

[AdrienCorenflos]

This is exactly what I had in mind indeed. I am not "concerned" about anything really, but it indeed has implications both in terms of code aesthetics and compilation times.

An example that breaks aesthetics would be:

What I'd like to be able to do

def f(x, theta):
    return theta * x

@jit
def g(h, y):
    return h(y)

g(jit(lambda x: f(x, 1.)), 0.)  # would be good if we could "partial" a jit without breaking jitting, but that's another point
g(jit(lambda x: f(x, 2.)), 0.)

What is feasible

@jit
def f(x, theta):
    return theta * x

@partial(jit, static_argnums=0)
def g(h, y, h_params):
    return h(y, *h_params)

g(f, 1., 1.) 
g(f, 1., 2.) 

On the side of the compilation speed, a good example is the fixed point routine example. What happens if I need to call this for a number of functions? While loops are pretty slow to compile.

Either way you did answer my somewhat academic question, and as I understand it it's a matter of static graph optimization (which is XLA promise)