Using functools.cached_property

[mjo22]

Hello! I had a question about how Modules work under jit with the functools.cached_property decorator.

I recently came across this decorator and thought it was pretty nifty. I have several quantities I want to compute in my Modules that depend on the leaves, and I’d like to use this solution in certain cases rather than writing a custom __init__.

Is it understood how cached properties work under jit? Thanks so much.