DESC icon indicating copy to clipboard operation
DESC copied to clipboard

Don't close over `self` when jitting objective functions

Open f0uriest opened this issue 1 year ago • 3 comments

  • This is a source of bugs, since attributes of the class get baked in, so don't get properly updated if they change
  • This also likely makes compilation slower, since it's trying to bake in some big constants
  • Also likely a source of memory leaks

f0uriest avatar Mar 27 '24 16:03 f0uriest

This issue might be related to a recent commit 9a78600 where I found I had to use np instead of jnp when working with class attributes that should be constant in jitted class functions.

ddudt avatar Mar 27 '24 16:03 ddudt

Could add a unique ID when a DESC object is created (during its init), to fix some issues that Rory was running into when trying to change this

dpanici avatar May 08 '24 20:05 dpanici

So one option for a unique identifier could be adding a __new__ method to the IOAble base class:

    def __new__(cls):
        self = object.__new__(cls)
        self._id = id(self)
        return self

so that _id is now saved as part of the objects __dict__, so it will be the same after jax tree flattening/unflattening.

The issue with this is that the _id is also the same after doing newthing = oldthing.copy() which we probably won't want.

Another option would be to use the existing thing.equals(otherthing) method? Though that could break if for some reason there are 2 objects with the same values in an optimization? (ie coils with same initial parameters?)

f0uriest avatar May 08 '24 22:05 f0uriest