Use "real" python code inside of jit

[inversecrime]

My functions often compute values that do not depend on the contents of any of the inputs, but rather on the shape of the inputs or not on the inputs at all.
This means that theoretically, it would be possible to precompute these values when jitting these functions, since jit assumes that the inputs have fixed shapes (as far as i understand, calling a jitted function with inputs of different shapes just triggers another compilation).

Is there any way to stage these precomputations out of jit?
A simple example would be the following code:

from functools import cached_property

import jax
import jax.numpy as jnp
from jax import Array


class Foo:
    @cached_property
    def some_property(self) -> Array:
        return jnp.ones((10, 10))


foo = Foo()
foo.some_property # depending on whether or not this is called, either (1) or (2) happens
jax.jit(lambda: foo.some_property)()
print(foo.some_property)  # (1) returns 10x10 ones (2) returns a traced array

This behavior is exactly what we expect from jit.
However, I want to achieve result (1) without calling the property before the jit command.

[jakevdp]

It sounds like you're looking for jax.ensure_compile_time_eval:

class Foo:
    @cached_property
    def some_property(self) -> Array:
        with jax.ensure_compile_time_eval():
            return jnp.ones((10, 10))

[jakevdp]

I wonder if this question is a form of premature optimization – under JIT, the XLA compiler will de-duplicate repeated operations, so perhaps the best approach is to not worry about cacheing at all, and just let the compiler do its job.

[inversecrime]

The thing is, I want to solve PDEs, and usually, creating the "pde sparse matrix" is computationally expensive (not as much as the solver of course, but still relevant).
When solving multiple PDEs with the same matrix, the compiler does not seem to recognize that the matrix does not change from call to call, and the jitted function recomputes it every single time.

Using vmap works well for this (then the matrix only gets created once), but sometimes it is not possible to use it.

I guess the obvious solution would be to create the matrix in a seperate function and then only jit the solver.

[jakevdp]

Yeah, that sounds better than trying to do caching via side-effects.

[inversecrime]

I don't really understand why caching is a side-effect. Would it be better practise to do this:

class Foo:
    def __init__(self) -> None:
        self.some_property = jnp.ones((10, 10))

Than this:

class Foo:
    @cached_property
    def some_property(self) -> Array:
        with jax.ensure_compile_time_eval():
            return jnp.ones((10, 10))

?

Is calling some_property considered a side-effect in both cases?
From my understanding after reading https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#pure-functions, it would not be a side effect if you pass an object of Class Foo to the corresponding function, like this:

@register_static
class Foo:
    def __init__(self) -> None:
        self.some_property = jnp.ones((10, 10))


def f(foo: Foo, x: Array) -> Array:
    return foo.some_property * x


print(jax.jit(f)(Foo(), jnp.ones(10)))

[jakevdp]

Cacheing is a side-effect because it mutates the object in-place without returning the mutated object: in other words, as a side-effect of calling the function, the object's contents change.

In the second case, it's also a side-effect, but it's a side-effect that's carefully constructed so as to not fall afoul of the assumptions made by JAX's transformation infrastructure. It's probably not a perfect solution, which is why I suggested a more direct approach.

[patrick-kidger]

To add on to the answer here. Compiling will generally precompute those values that depend only on array shapes and known constants. However, there are some exceptions to this. One example is that a * 0 will not reduce to 0, because nan * 0 == nan, so that's not a valid optimisation to make. This has quite far-reaching consequences; for example it's the reason that Jacobian computations often have an extra matmul against the identity matrix. (Which does not perform an identity operation if the function has NaNs.)

Another exception is if the number of dimensions in the constants get too large. Past some point the constant propagation will judge that it's too expensive, and defer the computation to runtime.

---

Wrt the specifics of your problem here, indeed explicitly precomputing the desired matrix in advance is probably what you want.

By the way, if you're solving differential equations, then you might like Diffrax. It's mostly targeting ODE/SDE solving, but it can handle some semidiscretised PDEs as well.

[inversecrime]

Thanks for your reply.
Indeed, I'm mainly interested in the wave equation and solving it with semidiscrete methods. I'll have a look at Diffrax.
I guess there is no scheme for 2nd order odes? Of course, I could just represent it as 1st order ode.