flax
5a82a055 - Fix setup trigger issue under sharing and transforms.

Commit
2 years ago
Fix setup trigger issue under sharing and transforms. There was a subtle bug in the case of using a shared module that defined a setup() method that was called first at one location under a transform (like remat or jit) and at another location without the transform. In the latter case the "outside" instance of the submodule had had a "shallow" _try_setup call pre-transform and was then marked as having had setup called, but the setup method had not in fact been called. This requires a one-line logic fix in _try_setup().
Author
Committer
Parents
Loading