DESC
DESC copied to clipboard
Don't close over `self` when jitting objective functions
- 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
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.
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
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?)