Jitting a function that returns a function

[attila-i-szabo]

I'm writing a factory function that returns the function v -> JᵀJv where J is the Jacobian of some function. If I write

@partial(jit, static_argnums=0)
def factory(fn, primals):
    _, jvp_fn = linearize(fn, primals)
    _, vjp_fn = vjp(fn, primals)
    def mat_vec(v):
        return vjp_fn(jvp_fn(v))
    return mat_vec

I get the error

TypeError: Value <function factory.<locals>.mat_vec at 0x7ff9706ea940> with type <class 'function'> is not a valid JAX type

which I take to mean that functions cannot be returned from a jitted function.

My use case for this requires this matrix-vector product with many vectors (not known in advance) for the same Jacobian (so I want to return a function), but I also use several different primals, which is why I would like to jit the factory. Is there any way to achieve this? (The function I actually want to return is somewhat more complicated, so I'd like a solution for arbitrary functions that don't do any unJAXy things to primals.)

[jakevdp]

Thanks for the question! A jitted function cannot return a function, but an un-jitted function can return a jitted function:

def factory(fn, primals):
    _, jvp_fn = linearize(fn, primals)
    _, vjp_fn = vjp(fn, primals)
    @jit
    def mat_vec(v):
        return vjp_fn(jvp_fn(v))
    return mat_vec

would that be sufficient for your use case?

[attila-i-szabo]

Yep, I discovered this too, but jitting the factory makes too much of a difference. I think I'll opt for jitting bigger blocks that return ndarrays. Good to know for sure though!

[inailuig]

With #7101 you might be able to do something like:

import jax
from functools import partial


@partial(jax.jit, static_argnums=0)
def _fwd(fn, primals):
    _, jvp_fn = jax.linearize(fn, primals)
    _, vjp_fn = jax.vjp(fn, primals)
    return jvp_fn, vjp_fn

@jax.jit
def _mat_vec(jvp_fn, vjp_fn, v):
    return vjp_fn(jvp_fn(v))

def factory(fn, primals):
    jvp_fn, vjp_fn = _fwd(fn, primals)
    return partial(_mat_vec, jvp_fn, vjp_fn)

[mattjj]

I think we can also write it using Partial like this:

from functools import partial

import jax
from jax.tree_util import Partial

@partial(jax.jit, static_argnums=0)
def factory(fn, primals):
    _, jvp_fn = jax.linearize(fn, primals)
    _, vjp_fn = jax.vjp(fn, primals)
    def mat_vec(jvp_fn, vjp_fn, v):
        return vjp_fn(jvp_fn(v))
    return Partial(mat_vec, jvp_fn, vjp_fn)

gnvp = factory(lambda x: x ** 2, 1.)
print(gnvp(1.))

[mattjj]

I think to use jax.linear_transpose here, we'd need to do a similar Partial trick to its implementation (in jax/_src/api.py).

[PhilipVinc]

With simple enough functions jax does CSE correctly (at least, inspecting the HLO dump).
I'd like to convince myself that this also happens with more complex functions but the HLO gets hard to read.
@mattjj, how does Jax/XLA perform CSE? is this documented anywhere? is it guaranteed to do it perfectly or is there some heuristics used to avoid spending too much time compiling?

[mattjj]

CSE is all in XLA rather than JAX. The code is all open-source, but I don't know if it's documented anywhere.

For functionally pure computations not involving control flow (i.e. not involving while_loop and cond), CSE is pretty easy, and I expect XLA to do it essentially perfectly. (At least, it is when all function applications are inlined, but XLA does this anyway, for reasons other than CSE as I understand it. That's one reason XLA compile times can get long!) XLA doesn't do much or any CSE across while_loop or cond boundaries, though I think that's just a current implementation detail and not a guarantee.

I'm sure there's much more that can be said, but that's my basic understanding.

[inailuig]

following up on this I opened #7191, which allows linearizing only once without relying on CSE:

from functools import partial

import jax
import jax.numpy as jnp
from jax.tree_util import Partial

@partial(jax.jit, static_argnums=0)
def factory(fn, primals):
    _, jvp_fn = jax.linearize(fn, primals)
    def mat_vec(jvp_fn, v):
        vjp_fn = jax.linear_transpose(jvp_fn, v)
        return vjp_fn(jvp_fn(v))
    return Partial(mat_vec, jvp_fn)

A = jnp.array([[1.,2.],[3.,4.]])
x = jnp.array([5., 6.])

def f(x):
    return A@x

gnvp = factory(f, x)

assert jnp.allclose(gnvp(x), A.T@A@x)
print(jax.jit(gnvp)(x))