In JAX, I changed a constant variable from a class member to a function argument, observing a different behavior. Here is the detail: I have a function in JAX that has a class object Obj as an input. So the execution of the function is something like func(Obj, ...), where ... is some other inputs. I had a constant variable as a member of a class, like Obj.var. I used pytree to handle the class, as suggested in JAX documentation. I made Obj.var static when registering the pytree because I don't need to change it. Recently, I removed var from the class and made it an argument of a function, like f(Obj, var, ...). I didn't make var static because I may want to batch over it. Then, the result of the function changes. Is this something expected? What is the good practice to handle a constant value like this? Continue reading...
