diff --git a/numpyro/distributions/transforms.py b/numpyro/distributions/transforms.py index d2e332bdc..0835d6a77 100644 --- a/numpyro/distributions/transforms.py +++ b/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 diff --git a/test/test_distributions.py b/test/test_distributions.py index 5533e7287..cd6cb9c42 100644 --- a/test/test_distributions.py +++ b/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