Class with jitted function: the following solution is it a good practice?

[jecampagne]

Hello,

I know that using jit for class methods should be done with care to avoid side effects. May be one can point to me a link that explains some use case/good practice in this subject (for instance I have read issue #1567 where @mattjj gives some interesting ways, and I encourage to set a specific documentation if not already done).

Well, I would like to submit to your comment/critics the following simple code that seems to work but may be more experienced jax-developpers would like to improve it. Notice that in my real-use case I have a tenth of "a" and "b" variables... The "a" variable is an archetype of variable that will not change during the processing of function like the "f"-func, while "b" variable is an archetype of variable that will change during the processing.

So now you know the spirit/context, here is the snippet:

import jax
import jax.numpy as jnp
from jax import vmap, jit
from functools import partial
jax.config.update("jax_enable_x64", True)


class A():
    def __init__(self, a: jnp.array)->None:
        self.a = a
        self.Init()
        
    def Init(self)->None:
        self.b = None
        
    def set_b(self, x):
        self.b = x
    
    @partial(jit, static_argnums=(0,))
    def f(self, var: float)->float:
        b = self.a * var
        return b

objA = A(jnp.array([2.0]))
print("1)",objA.a, objA.b)

b = objA.f(10.)
print("2)",b)
objA.set_b(b)
print("3)",objA.a, objA.b)

new_objA = A(jnp.array([3.0]))
print("4)",objA.a, objA.b)
print("5)",new_objA.a, new_objA.b)

new_b = new_objA.f(20.)
new_objA.set_b(new_b)
print("6)",objA.a, objA.b)
print("7)",new_objA.a, new_objA.b)

This gives as expected

1) [2.] None
2) [20.]
3) [2.] [20.]
4) [2.] [20.]
5) [3.] None
6) [2.] [20.]
7) [3.] [60.]

So, now 1-2-3 start to comment... Thanks.

[YouJiacheng]

I think the best practice is to make your class a pytree, and use class method instead of instance method.
IIUC, your solution will suffer excessive re-compilation and might silently cause error.

[jecampagne]

Great !

1. can you propose a better solution
2. can you elaborate a little mode on "your solution will suffer excessive re-compilation".
   Thanks.

[YouJiacheng]

@jecampagne

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


class A():
    def __init__(self, a: jnp.ndarray):
        self.a = a
        self.b = jnp.zeros_like(a)

    def set_b(self, x):
        self.b = x

    @partial(jit, static_argnums=(0,))
    def f(self, var: float) -> float:
        print('compiling')
        b = self.a * var + self.b
        return b


objA = A(jnp.array([2.0]))
print("1)", objA.a, objA.b)

b = objA.f(10.)
print("2)", b)
objA.set_b(b)
print("3)", objA.a, objA.b)
print("3.1)", objA.f(10.)) # expect 40.

new_objA = A(jnp.array([3.0]))
print("4)", objA.a, objA.b)
print("5)", new_objA.a, new_objA.b)

new_b = new_objA.f(20.) # expect not re-compiling
new_objA.set_b(new_b)
print("6)", objA.a, objA.b)
print("7)", new_objA.a, new_objA.b)

# 1) [2.] [0.]
# compiling
# 2) [20.]
# 3) [2.] [20.]
# 3.1) [20.]
# 4) [2.] [20.]
# 5) [3.] [0.]
# compiling
# 6) [2.] [20.]
# 7) [3.] [60.]

[jakevdp]

Marking self as static can be dangerous if you've not defined an appropriate __hash__ function for your type. Here's a way that this code can go wrong:

obj = A(1)
print(obj.f(2))
# 2

obj.a = 2
print(obj.f(2))  # should print 4, but prints 2
# 2

The issue is that static arguments to jit are evaluated based on their hash, but your object's hash does not take into account the value of self.a. Thus at the second function call, the hash has not changed, so JAX's machinery assumes the value of self.a in the function has not changed.

However, making __hash__ respect the value self.a can lead to excessive recompilation – each time self.a changes, the jit-compiled method will have to be re-compiled.

For this reason, I would instead make your class a custom pytree and no longer mark the self argument as static. Here is an example:

class A():
    def __init__(self, a: jnp.array)->None:
        self.a = a
        self.Init()
        
    def Init(self)->None:
        self.b = None
        
    def set_b(self, x):
        self.b = x
    
    @jit
    def f(self, var: float)->float:
        b = self.a * var
        return b

    def _tree_flatten(self):
      # You might also want to store self.b in either the first group
      # (if it's not hashable) or the second group (if it's hashable)
      return (self.a,), ()

    @classmethod
    def _tree_unflatten(cls, aux, children):
      return cls(*children)

tree_util.register_pytree_node(A, A._tree_flatten, A._tree_unflatten)

obj = A(1)
print(obj.f(2))
# 2

obj.a = 2
print(obj.f(2))
# 4

Note that I ignored the value of self.b here, because it's unclear to me from the class design what its intent is. You can return it in the children if it's meant to be an (unhashable) array, and in aux_data if it's meant to be a hashable parameter.

[jakevdp]

Yes, but that's playing a dangerous game: if you ever pass your object to a JAX transform and use an attribute that's not represented by the tree flattening, even indirectly, it will essentially be treated as a compile-time constant and subsequent function calls may produce unexpected results (unexpected, but correct according to how you declared the object should behave!)

[moehmeni]

Can we use register_pytree_node inside the class itself (e.g. in init) so the user does not need to do that when using the class? also can the class be dataclass like flax modules?
One more thing is that when I test initilixing the jit method in __init__, it is faster than the decorating it with @jax.jit for the same method:
fast:

@register_pytree_node_class
class Test:
    def __init__(self):
        self.method_jit = jax.jit(self.method)
        key = jax.random.PRNGKey(0)
        self.a = jax.random.normal(key, (1000, 1000))
        self.b = jax.random.normal(key, (1000, 1000))
        self.method_jit_init = jax.jit(self.method)

    @jax.jit
    def method_jit(self):
        return jnp.dot(self.a, self.b)

    def method(self):
        return jnp.dot(self.a, self.b)

    def tree_flatten(self):
        return (), None

    @classmethod
    def tree_unflatten(cls, aux_data, children):
        return cls(*children)

t = Test()
fns = [t.method, t.method_jit, t.method_jit_init]
for fn in fns:
    times = []
    for _ in range(100):
        t1 = time.time()
        fn()
        t2 = time.time()
        times.append(t2 - t1)
    print(f"{fn.__name__}: {sum(times) / len(times)}")

method: 0.004103074073791504
method: 0.0004682230949401855
method: 0.0004332923889160156

is there any advantage for declaring the function in __init__?

[jakevdp]

One issue with this is that as currently written, if you later modify t.a or t.b, those modified attributes will not be updated when you call method_jit.

[moehmeni]

Thanks I realized what was the problem, but I noticed when I try to do the same thing for flax inherited modules:

@register_pytree_node_class
@dataclass
class Parent(nn.Module):
    key: jax.random.PRNGKey
    params: dict = None

    @jax.jit
    def step(self, loss_fn, optimizer, opt_state, x, y):
       # ...

    @jax.jit
    def predict(self, x):
        # ...

    def fit( # ... ):
            # ...
            for x, y in utils.data_loader(X, Y, batch_size):
                self.params, opt_state, loss_value = self.step(
                    loss, optimizer, opt_state, x, y
                ) # I get the error on this line
            # ...

    def tree_flatten(self):
        return (self.params,), None

    @classmethod
    def tree_unflatten(cls, aux_data, children):
        return cls(*children, aux_data)

class TestModel(Parent):
    d_hidden: int = 64
    d_out: int = 1

    @nn.compact
    def __call__(self, x):
        # ...

I get this error when calling fit() of the child class.

TypeError: Cannot interpret value of type <class '__main__.TestModel'> as an abstract array; it does not have a dtype attribute

Is there something missing? should not register_pytree_node_class solve this problem?

[jakevdp]

You would get that error if you're passing the Module object to a JIT function. It's hard to know what to recommend without seeing more of your code, but a deep comment thread on a multi-year-old question is not the best forum for that.

I'd suggest asking your question at https://github.com/google/flax/discussions

[YouJiacheng]

Best practice IMO:
Oh, jakevdp's solution is definitely better.

import jax.numpy as jnp
from jax import jit
from jax.tree_util import register_pytree_node
from functools import partial


class A(object):
    def __init__(self, a, b):
        self.a = a
        self.b = b

    @classmethod # classmethod is a descriptor, not a callable
    @partial(jit, static_argnums=(0,))
    def f(cls, obj: 'A', v):
        print('compiling')
        return obj.a * v + obj.b

register_pytree_node(
    A,
    lambda objA: ((objA.a, objA.b), None),
    lambda _, children: A(*children)
)

objA = A(jnp.array(1.0), jnp.array(5.0))
print(objA.f(objA, 10.0))
objA.b = 10.0
print(objA.f(objA, 10.0))
new_objA = A(jnp.array(6.0), jnp.array(5.0))
print(new_objA.f(new_objA, 10.0))

# compiling
# 15.0
# 20.0
# 65.0

[jecampagne]

Once again many thanks for your help. Even if custom pytree gives some valuable instructions, I think sometimes basic example also help beginners to understand the need to use at-first-glance complicated structures.

[jakevdp]

See https://jax.readthedocs.io/en/latest/faq.html#how-to-use-jit-with-methods

[jecampagne]

Super cool ! Thanks

[patrick-kidger]

Much easier for nearly all use cases is to use a libary that treats classes as PyTrees. I'd recommend Equinox. (Of which I am the author.) This will allow you to use classes/jax.jit/etc. together in pretty much arbitrary ways.

[jecampagne]

Hi @patrick-kidger

It would be nice if you can write your way to code the "A-class". note that "a" variables are archetype of variable that is not changed after creation of an A-object, while "b" variables would change/be set by functions, the "f-func" is an archetype of jit-function that use "a" variable and return "b"-type var. Finally (even not mentionned) there are possibly "g"-type functions that cannot be jitted (eg. because of use of third party lib, as jaxopt.ScipyBoundedMinimize for instance).

So, let us know how you would proceed? Thanks

[jecampagne]

Hi, following #1567 I have cooked this snippet
but I wander why it seems to indicate a recompilation between print 10 and 11 using both jnp.array while there is no sign of compilation using float (form print 2 to 9).

An idea?

clear_cache = jax._src.dispatch._xla_callable.cache_clear

def u_print(idx:int, *args)->int:
    print(f"{idx}):",*args)
    idx +=1
    return idx

########

class A():
    def __init__(self, a, b=None):
#        print("Nouveau A")
        self.a = a
        self.b = b

    @jit                                       
    def f(self, var):
        print("compile...")
        new_b = self.a * var
        return A(self.a, new_b)       # on retourne un nouvel objet.


tree_util.register_pytree_node(A,
                               lambda x: ((x.a,x.b), None),
                               lambda _, x: A(a=x[0],b=x[1]) 
                              )

###############

import time

clear_cache
idp=0

objA = A(2.0)
idp = u_print(idp, objA.a, objA.b)    #0

objA = objA.f(10.)
time.sleep(0.5)
idp = u_print(idp,objA.a, objA.b)     #1

objA = objA.f(11.)
time.sleep(0.5)
idp = u_print(idp, objA.a, objA.b)    #2

####

objA = A(3.0) 
idp = u_print(idp, objA.a, objA.b)     #3
objA = objA.f(20.)
time.sleep(0.5)
idp = u_print(idp, objA.a, objA.b)     #4
objA= objA.f(30.)
time.sleep(0.5)
idp = u_print(idp, objA.a, objA.b)     #5


objA.a = 30.  # permis mais volontaire
objA= objA.f(30.)
time.sleep(0.5)
idp = u_print(idp, objA.a, objA.b)     #6

####

objA = A(3.0) 
idp = u_print(idp, objA.a, objA.b)     #7

objA= objA.f(20.)
time.sleep(0.5)
idp = u_print(idp, objA.a, objA.b)     #8

objA= objA.f(30.)
time.sleep(0.5)
idp = u_print(idp, objA.a, objA.b)     #9

print(">>>> try an array for var")

objA= objA.f(jnp.array([30.]))
time.sleep(0.5)
idp = u_print(idp, objA.a, objA.b)     #10
time.sleep(0.5)

objA= objA.f(jnp.array([40.]))
time.sleep(0.5)
idp = u_print(idp, objA.a, objA.b)     #11

#objA.a = 40. ###
#objA= objA.f(jnp.array([40.]))
#time.sleep(0.5)
#idp = u_print(idp, objA.a, objA.b)     #12

leading to

0): 2.0 None
compile...
1): 2.0 20.0
compile...
2): 2.0 22.0
3): 3.0 None
4): 3.0 60.0
5): 3.0 90.0
6): 30.0 900.0
7): 3.0 None
8): 3.0 60.0
9): 3.0 90.0
>>>> try an array for var
compile...
10): 3.0 [90.]
compile...
11): 3.0 [120.]

[YouJiacheng]

for the f call between 9 and 10, b is 90.0, as shown in 9.
for the f call between 10 and 11, b is [90.0], as shown in 10, which is different from the call between 9 and 10, thus recompiled.
you can use jnp.array(30.0) instead of jnp.array([30.0]).

[jecampagne]

I think the way is the following

• for 9) a = 3.0 and arg of f is 30. -> b=90. but as in 8) there is only flaot => no compilation needed
• for 10) a= 3.0 and arg of f is [30.] -> b=[90.] but as f was using (float,float) => need recompilation to get (float, jnp.array[x])
  -for 11) a=3.0 and arg of f is [40.] -> b=[120] but self.b as changed from float to jnp.array[x] => need recompil.

This analysis is corroborated by the fact that if I add

objA= objA.f(jnp.array([50.]))         
time.sleep(0.5)
idp = u_print(idp, objA.a, objA.b)     #12

then no recompilation is triggered...

[jakevdp]

Remember that self is an implicit argument to every method, so if the content of self changes between function calls, it will lead to reconciliation.

[jecampagne]

Hi again,
A long the line to implement an helper solution to the JIT-in-class task, I face the following problem if the f-func needs an external fonction (ie. the g-func

def g(x):
     return x**2


class A():
    def __init__(self, a):
        print("Nouveau A")
        self.a = a
        self.g = g
        self.b = None

    def f(self, var):
        
        kargs={}
        kargs["var"] = var
        kargs["g"] = self.g
        
        self.b = _f(self.a,kargs)

        
@jit
def _f(a, kargs):
    print("compile...")
    var =  kargs["var"]
    g   =  kargs["g"]
    res = g(a)*var
    return res

clear_cache
idp=0

objA = A(2.0)
idp = u_print(idp, objA.a, objA.b)
objA.f(10.)

I get:

Nouveau A
0): 2.0 None

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Input In [53], in <cell line: 34>()
     32 objA = A(2.0)
     33 idp = u_print(idp, objA.a, objA.b)
---> 34 objA.f(10.)

Input In [53], in A.f(self, var)
     15 kargs["var"] = var
     16 kargs["g"] = self.g
---> 18 self.b = _f(self.a,kargs)

    [... skipping hidden 3 frame]

File /jax/_src/api.py:2687, in _check_arg(arg)
   2685 def _check_arg(arg):
   2686   if not (isinstance(arg, core.Tracer) or _valid_jaxtype(arg)):
-> 2687     raise TypeError(f"Argument '{arg}' of type {type(arg)} is not a valid JAX type.")

TypeError: Argument '<function g at 0x7fe92c4f9ee0>' of type <class 'function'> is not a valid JAX type.

Is there a solution considering the use of helper function? Thanks

[YouJiacheng]

Use jax.tree_util.Partial(g) to pass it as a pytree.

[jakevdp]

You can also mark the g argument as static in your jit compilation.

By the way, a new answer to an existing question is not a great place to ask a new question. In the future please feel free to open a new discussion (with all the unthreaded replies here, this one is getting to be pretty hard to follow, and the original question/answer and related discussion is pretty obscured.)

[jecampagne]

Thanks for both ot you.