Update MVN.expand() to support non-lazy MVN & reuse scale_tril where possible #2623
Conversation
| else: | ||
| # Non-lazy MVN is represented using scale_tril in PyTorch. | ||
| # Constructing it from scale_tril will avoid unnecessary computation. | ||
| # Initialize using __new__, so that we can skip __init__ and use scale_tril. |
There was a problem hiding this comment.
This seems slightly hairy - would an alternative here be to just add support for MultivariateNormal to be initialized directly from the scale_tril?
There was a problem hiding this comment.
Yeah, that'd be an option. I didn't want to touch it since it can potentially affect correctness of MVN. But it doesn't seem like _covar is directly used anywhere, so it may be safer than I initially though.
There was a problem hiding this comment.
Actually, what would we do if both covar and scale_tril are provided? Do we validate that they are compatible? This would require re-computing one or the other, not the cheapest operation. Seems to conflict with the goal of re-using whenever these are available.
There was a problem hiding this comment.
I feel like we could also just make it such that we error out if both of them are provided.
The previous implementation of
MVN.expand()would error out when MVN was not lazy (sinceself._covaris not defined) and would discard thescale_tril, requiring its re-computation down the line. This PR updatesMVN.expand()to support both lazy and non-lazy MVN and to re-usescale_trilwhenever possible.Test Plan:
Added unittests for
MVN.expand().