Fix incorrect unflattenning of inverse transforms (#1600)

pyro-ppl · Jun 6, 2023 · 523162f · 523162f
1 parent 4e3c50e
commit 523162f
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/numpyro/distributions/transforms.py b/numpyro/distributions/transforms.py
@@ -163,9 +163,10 @@ def inverse_shape(self, shape):
     def tree_flatten(self):
         return (self._inv,), (("_inv",), dict())
 
-    @classmethod
-    def tree_unflatten(cls, aux_data, params):
-        return cls(params)
+    def __eq__(self, other):
+        if not isinstance(other, _InverseTransform):
+            return False
+        return self._inv == other._inv
 
 
 class AbsTransform(ParameterFreeTransform):

diff --git a/test/test_transforms.py b/test/test_transforms.py
@@ -140,6 +140,8 @@ def out_t(transform, x):
     assert jitted_in_t(transform, 1.0) == 1.0
     assert jitted_out_t(transform, 1.0) == transform
 
+    assert jitted_out_t(transform.inv, 1.0) == transform.inv
+
     assert jnp.allclose(
         vmap(in_t, in_axes=(None, 0), out_axes=0)(transform, jnp.ones(3)),
         jnp.ones(3),