Partial Evaluation / Precomputation

[inversecrime]

I just had the idea for this function, "precompute", which basically uses jax.interpreters.partial_eval to precompute all values associated with certain inputs.

from jax import make_jaxpr
from jax import numpy as jnp
from jax.core import eval_jaxpr
from jax.interpreters.partial_eval import partial_eval_jaxpr_nounits


def precompute(func, arg_list, known_argnums):
    jaxpr = make_jaxpr(func)(*arg_list)
    unknowns = [False if k in known_argnums else True for k in range(len(arg_list))]
    instantiate = False

    (known_jaxpr, unknown_jaxpr, _, _) = partial_eval_jaxpr_nounits(jaxpr, unknowns, instantiate)

    known_values = [arg_k for (k, arg_k) in enumerate(arg_list) if k in known_argnums]
    precomputed_values = eval_jaxpr(known_jaxpr.jaxpr, known_jaxpr.consts, *known_values)

    def inner(*args):
        unknown_values = [arg_k for (k, arg_k) in enumerate(args) if k not in known_argnums]
        values = eval_jaxpr(unknown_jaxpr.jaxpr, unknown_jaxpr.consts, *precomputed_values, *unknown_values)
        return values
    return inner


def f(x, y, z):
    def g(y):  # Some complicated calculation based on y
        return (2 + jnp.ones((10, 10)))**2
    return (x * g(y) + z)


y = 2 * jnp.ones(10)
z = jnp.ones_like(y) * 1.5
x = jnp.zeros_like(y)

f_precomputed = precompute(f, [jnp.zeros_like(x), y, jnp.zeros_like(z)], known_argnums=(1,))

# Both give the same (equivalent) jaxpr.
print(make_jaxpr(f)(x, y, z))
print(make_jaxpr(lambda x, y, z: precompute(f, [x, y, z], known_argnums=(1,))(x, y, z))(x, y, z))

# Values associated with y are not computed, but treated as constants. The input y is ignored, so both give the same jaxpr.
print(make_jaxpr(f_precomputed)(x, y, z))
print(make_jaxpr(f_precomputed)(x, 123, z))

The code seems kind of clunky, but in my example, it works as intended. Is there a chance that in some scenario, it does something unexpected?
Also, does something like this already exist?