Slow compilation when jitting method of a class

[IrishWhiskey]

I get the error E external/xla/xla/service/slow_operation_alarm.cc:65] Constant folding an instruction is taking > 1s when jitting the method of a class. Is that normal? Does JAX support jitting methods?

Here is the code that produces the error

import numpy as np
import jax
import jax.numpy as jnp
from functools import partial

np.random.seed(0)

class C:
    def __init__(self, d, w):
        self.d = d
        self.w = w
    
    @partial(jax.jit, static_argnums=(0,))
    def foobar(self):
        return self.d @ self.w
    
d = jnp.array(np.random.rand(200000, 3))
w = jnp.array(np.random.rand(3,))

c = C(d, w)
c.foobar()

I'm using JAX 0.4.10, Cuda 11.7, Cudnn 8.6 and Python 3.10.

[jakevdp]

The issue is that you are marking self as static. This is problematic for a number of reasons, as described in FAQ: How to use JIT with Methods.

In particular, if self is marked static then any attributes of self will be embedded as literal constants in the compiled program, and embedding a 600,000-element constant is going to lead to poor performance (thus the warning you're seeing).

I would recommend following Option 3 at the above link; it might look something like this:

import numpy as np
import jax
import jax.numpy as jnp
from functools import partial

np.random.seed(0)

@jax.tree_util.register_pytree_node_class
class C:
    def __init__(self, d, w):
        self.d = d
        self.w = w
    
    @jax.jit
    def foobar(self):
        return self.d @ self.w

    def tree_flatten(self):
        children = (self.d, self.w)  # arrays / dynamic values
        aux_data = {}  # static values
        return (children, aux_data)

    @classmethod
    def tree_unflatten(cls, aux_data, children):
        return cls(*children, **aux_data)
    
d = jnp.array(np.random.rand(200000, 3))
w = jnp.array(np.random.rand(3,))

c = C(d, w)
c.foobar()

If you run this you'll find that it compiles more-or-less instantly.

[IrishWhiskey]

Thank you :)

[maurorigo]

@jakevdp what about attributes that are not passed when initializing the class (but are computed inside __init__) or even that are defined inside some other class method? I can't really find documentation on how to handle these

[jakevdp]

You can reconstruct your object manually using object.__new__ if doing so via __init__ doesn't make sense; there is a brief example at https://jax.readthedocs.io/en/latest/pytrees.html#custom-pytrees-and-initialization, but the pattern you mention seems to be based on in-place mutation of the object (if calling methods is modifying object attributes), so there are other issues you should worry about as well, documented elsewhere on that page.

[maurorigo]

Thank you very much!

[patrick-kidger]

Follow-up: if you want a nicer syntax for Jake's answer, then try Equinox. This defines equinox.Module, which automatically handles the pytree registration, and many other edge cases, for you.