Memoize Transform.inv (#885)

Author: fritzo

numpyro/distributions/transforms.py

@@ -3,6 +3,7 @@
 
 import math
 import warnings
+import weakref
 
 import numpy as np
 
@@ -52,6 +53,7 @@ def _clipped_expit(x):
 class Transform(object):
     domain = constraints.real
     codomain = constraints.real
+    _inv = None
 
     @property
     def event_dim(self):
@@ -62,7 +64,13 @@ def event_dim(self):
 
     @property
     def inv(self):
-        return _InverseTransform(self)
+        inv = None
+        if self._inv is not None:
+            inv = self._inv()
+        if inv is None:
+            inv = _InverseTransform(self)
+            self._inv = weakref.ref(inv)
+        return inv
 
     def __call__(self, x):
         return NotImplementedError

test/test_distributions.py

@@ -1038,6 +1038,7 @@ def test_bijective_transforms(transform, event_shape, batch_shape):
     z = transform.inv(y)
     assert_allclose(x, z, atol=1e-6, rtol=1e-6)
     assert transform.inv.inv is transform
+    assert transform.inv is transform.inv
     assert transform.domain is transform.inv.codomain
     assert transform.codomain is transform.inv.domain