Problem in jitting methods of a class

[mtagliazucchi]

Hi, I'm trying to JIT a methods of a class that updates the properties of that class.
Here's a minimal working example:

import jax
import jax.numpy as jnp

@jax.jit
def _simple_function(x, par1, par2):
    return par1 + par2*x

class ExampleClass():
    def __init__(self, Lambda=None, vectorize=False):
        
        fidLambda = {'par1':1.0, 'par2':2.0}
        self.keys = fidLambda.keys()
        
        self.Lambda = Lambda if Lambda is not None else fidLambda
        self.fidLambda = self.Lambda.copy() 
        self.vectorize = vectorize

    @jax.jit  
    def get_ordered_params(self):
        return list(self.Lambda[key] for key in self.keys)
    
    @jax.jit
    def update_params(self, Lambda=None):
        
        if Lambda is None:
            self.vectorize_computation = False
            return self

        invalid_keys = [k for k in Lambda.keys() if k not in self.keys]
        if invalid_keys:
            raise ValueError(f"`Lambda` contains invalid keys {invalid_keys}. Valid keys are {self.keys}.")   
            
        up_params            = {k:Lambda[k] for k in Lambda}
        sizes_updated_params = [Lambda[k].size for k in Lambda]
        
        if max(sizes_updated_params) == 1:
            self.vectorize_computation=False
            self.Lambda.update(up_params) 
            self.Lambda.update({k:self.fidLambda[k] for k in self.keys if k not in up_params})
        else:
            self.vectorize_computation=True
            non_up_params = {k:jnp.full(max(sizes_updated_params), self.fidLambda[k]) for k in self.keys if k not in up_params}
            self.Lambda.update(up_params)
            self.Lambda.update(non_up_params)
        return self
    
    @jax.jit
    def simple_function(self, x):
        if self.vectorize:
            return vmap(_simple_function, in_axes=(None,0,0))(x, *self.get_ordered_params())
        else:
            return _simple_function(x, *self.get_ordered_params())
        
    def _tree_flatten(self):
        children = (self.Lambda, )  
        aux_data = {'vectorize':self.vectorize} 
        return (children, aux_data)

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

jax.tree_util.register_pytree_node(ExampleClass, ExampleClass._tree_flatten, ExampleClass._tree_unflatten)

This class is almost okay, but I noticed that the method update_params returns a new instance of the class with the updated parameters and it does not modify the parameters of the original instance. Here's the code to reproduce this behavior:

obj = ExampleClass()
print(obj.Lambda) # {'par1':1.0, 'par2':2.0}
obj.update_params({'par1':2.0}) 
print(obj.Lambda) # should print {'par1':2.0, 'par2':2.0}, but print same dictionary as before

# The methods return a new istance of the class with the update params
obj2 = obj.update_params({'par1':2.0})
print(obj2.Lambda) # {'par1':2.0, 'par2':2.0} 

I would like to avoid to generate a new class instance every time I update the class parameters because it doesn't seem efficient to me. Removing the jit decorator in update_params solves the problem, but unfortunately I need to JIT that method since it's called from another jitted function.

Do you have any suggestion that can help me?

[jakevdp]

Thanks for the question – we have a FAQ entry on this topic here, and it looks like you're already using the pytree approach, which is great.

This class is almost okay, but I noticed that the method update_params returns a new instance of the class with the updated parameters and it does not modify the parameters of the original instance.

This is working as expected: JIT-compiled functions must be pure, and cannot operate via side-effects. By design, the input to the function will not be mutated as a result of executing the function because this is a side-effect (see JAX sharp bits: pure functions for some discussion of this).

I would like to avoid to generate a new class instance every time I update the class parameters because it doesn't seem efficient to me. Removing the jit decorator in update_params solves the problem, but unfortunately I need to JIT that method since it's called from another jitted function.

Your options are either to remove the JIT compilation to use in-place mutation, or to keep the JIT compilation and return the modified class. You can't have both JIT and mutation. That said, if you're calling update_params within another JIT-compiled function, you can typically rely on the compiler to avoid inefficient copies at runtime.

[mtagliazucchi]

Thanks you so much! I will manage the new istances and keep the JIT compilation!