From bdf6d77bfa8306129bb5c4c5437df6b0038a3b06 Mon Sep 17 00:00:00 2001 From: Fritz Obermeyer Date: Thu, 21 Jan 2021 14:29:06 -0500 Subject: [PATCH] Simplify --- numpyro/distributions/transforms.py | 10 +++++++++- test/test_distributions.py | 1 + 2 files changed, 10 insertions(+), 1 deletion(-) 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