Repository: pyro-ppl/numpyro

Issue:
Constraints object are not robust to pickling
Hi, and first of all, thank you for this awesome project!

Properties of `Transforms` object such as `codomain` rely on identity checks such as this one:
https://github.com/pyro-ppl/numpyro/blob/34e0cdf4fa0ab9a0300a0d894d6758419fb46f40/numpyro/distributions/transforms.py#L173
which are not robust to serialization roundtrips (of the form `pickle.loads(pickle.dumps(obj))`, which lead to the creation of new `Constraints` objects). Because of this, the following code snippet fails with: 

```python
In [1]: import jax.numpy as jnp
   ...: import pickle
   ...: from numpyro.distributions.transforms import AffineTransform
   ...:
   ...: a = AffineTransform(loc=jnp.array([0., 0.]), scale=jnp.array([1., 1.]))
   ...: print(a.inv._inv.codomain)
   ...:
   ...: a_roundtripped = pickle.loads(pickle.dumps(a))
   ...: print(a_roundtripped.inv._inv.codomain)
```

```python
Real()
---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
<ipython-input-2-399a4b34ce15> in <module>
      7
      8 a_roundtripped = pickle.loads(pickle.dumps(a))
----> 9 print(a_roundtripped.inv._inv.codomain)

~/repos/numpyro/numpyro/distributions/transforms.py in codomain(self)
    195                 )
    196         else:
--> 197             raise NotImplementedError
    198
    199     def __call__(self, x):

NotImplementedError:
```

There are at least 2 ways that I can think of to fix this issue: 

- either rely on equality checks instead of identity checks, and include sensible `__eq__` methods on constraints objects
- or reduce `constraints` objects using the little known singleton reducing semantics by making `Constraints.__reduce__` return a string, and not the traditional reconstruction tuple:
```python
In [1]: import pickle
   ...: class _Real:
   ...:     def __reduce__(self):
   ...:         # ask `pickle.load` to grab the top-level `real` object from the module in which _Real is defined
   ...:         return "real"
   ...: real = _Real()
   ...: pickle.loads(pickle.dumps(real)) is real
Out[1]: True
```
The latter option would be an elegant solution, only if it makes sense to treat constraints as singleton instances.


Discussion:
Interesting! Does this happen at many places? I think we can replace `if ... is constraints ...` by `if isinstance(...)` statement, e.g.
```
if isinstance(self.domain, type(constraints.real)):
```
> Interesting! Does this happen at many places? 

Looking for identity statements on `constraints` objects using `git grep 'is constraints'` on the `numpyro` root directory gives (discarding results from `tests_*` files):
```
numpyro/distributions/transforms.py:        if self.domain is constraints.real:
numpyro/distributions/transforms.py:        if self.domain is constraints.real:
numpyro/distributions/transforms.py:    if constraint is constraints.positive:
numpyro/distributions/transforms.py:    if constraint is constraints.unit_interval:
numpyro/distributions/truncated.py:            base_dist.support is constraints.real
numpyro/distributions/truncated.py:            base_dist.support is constraints.real
numpyro/distributions/truncated.py:            base_dist.support is constraints.real
numpyro/infer/autoguide.py:                if site["fn"].support is constraints.real or (
numpyro/infer/autoguide.py:                    and site["fn"].support.base_constraint is constraints.real
numpyro/infer/reparam.py:        assert support is constraints.circular
numpyro/infer/util.py:        if support is constraints.circular:
```

Which of course is only a rough lower bound of such occurences. But I could not find much more than that using more complex introspecting queries.

> I think we can replace if ... is constraints ... by if isinstance(...) statement, e.g.

This goes against the rather intuitive feeling that `constraints.real` should be treated as a singleton - I suspect that would any people (reviewer or new contributor)  unfamiliar with that specific issue write some form of check involving `constraints.real` in the future, they will probably resort to identity semantics and "re-break" the code. In addition, systematic testing of the robustness of the `numpyro` codebase against serialization roundtrips is likely to be difficult to achieve, and thus such mistakes might not be caught by tests. Thus, I'd advocate for keeping `constraints.real` compatible with identity semantics by reducing it as a singleton.
> This goes against the rather intuitive feeling that constraints.real should be treated as a singleton ...

Those are very good points! Yeah, I think having `__reduce__` method would be great. Would the solution apply for `positive` and `unit_interval`? I'm seeing that `positive` is an instance of `_GreaterThan` and `unit_interval` is an instance of `_Interval`. I guess we can add `if/else` statements to decide how to pickle the object? Something like
```
class _GreaterThan(Constraint):
    def __reduce__(self):
        if self.lower_bound is 0.0:
            return "positive"
        else:
            return (_GreaterThan, (self.lower_bound,))  # I'm not sure
```
    
> Would the solution apply for positive and unit_interval?

In my opinion, yes, and (unless I'm missing something) to all top-level `Constraint` instances of the `constraints` module.

> I guess we can add if/else statements to decide how to pickle the object?

As a POC this looks good. We can also directly target the objects of interests in the logic to prevent any risk of false positives using for instance `if self is positive`, or `if self is getattr(globals(), "positive")` checks (the second option prevents issues with forward referencing `positive`)

For the fallback logic, something along the lines of `return super(_GreaterThan, self).__reduce__()` should do the job :-) 
Sounds reasonable to me! Could you submit a PR for this, @pierreglaser?

Resolving patch:
diff --git a/numpyro/distributions/constraints.py b/numpyro/distributions/constraints.py
--- a/numpyro/distributions/constraints.py
+++ b/numpyro/distributions/constraints.py
@@ -95,7 +95,20 @@ def feasible_like(self, prototype):
         raise NotImplementedError
 
 
-class _Boolean(Constraint):
+class _SingletonConstraint(Constraint):
+    """
+    A constraint type which has only one canonical instance, like constraints.real,
+    and unlike constraints.interval.
+    """
+
+    def __new__(cls):
+        if (not hasattr(cls, "_instance")) or (type(cls._instance) is not cls):
+            # Do not use the singleton instance of a superclass of cls.
+            cls._instance = super(_SingletonConstraint, cls).__new__(cls)
+        return cls._instance
+
+
+class _Boolean(_SingletonConstraint):
     is_discrete = True
 
     def __call__(self, x):
@@ -105,7 +118,7 @@ def feasible_like(self, prototype):
         return jax.numpy.zeros_like(prototype)
 
 
-class _CorrCholesky(Constraint):
+class _CorrCholesky(_SingletonConstraint):
     event_dim = 2
 
     def __call__(self, x):
@@ -126,7 +139,7 @@ def feasible_like(self, prototype):
         )
 
 
-class _CorrMatrix(Constraint):
+class _CorrMatrix(_SingletonConstraint):
     event_dim = 2
 
     def __call__(self, x):
@@ -231,6 +244,11 @@ def feasible_like(self, prototype):
         return jax.numpy.broadcast_to(self.lower_bound + 1, jax.numpy.shape(prototype))
 
 
+class _Positive(_GreaterThan, _SingletonConstraint):
+    def __init__(self):
+        super().__init__(0.0)
+
+
 class _IndependentConstraint(Constraint):
     """
     Wraps a constraint by aggregating over ``reinterpreted_batch_ndims``-many
@@ -284,6 +302,16 @@ def feasible_like(self, prototype):
         return self.base_constraint.feasible_like(prototype)
 
 
+class _RealVector(_IndependentConstraint, _SingletonConstraint):
+    def __init__(self):
+        super().__init__(_Real(), 1)
+
+
+class _RealMatrix(_IndependentConstraint, _SingletonConstraint):
+    def __init__(self):
+        super().__init__(_Real(), 2)
+
+
 class _LessThan(Constraint):
     def __init__(self, upper_bound):
         self.upper_bound = upper_bound
@@ -339,6 +367,16 @@ def feasible_like(self, prototype):
         return jax.numpy.broadcast_to(self.lower_bound, jax.numpy.shape(prototype))
 
 
+class _IntegerPositive(_IntegerGreaterThan, _SingletonConstraint):
+    def __init__(self):
+        super().__init__(1)
+
+
+class _IntegerNonnegative(_IntegerGreaterThan, _SingletonConstraint):
+    def __init__(self):
+        super().__init__(0)
+
+
 class _Interval(Constraint):
     def __init__(self, lower_bound, upper_bound):
         self.lower_bound = lower_bound
@@ -367,6 +405,16 @@ def __eq__(self, other):
         )
 
 
+class _Circular(_Interval, _SingletonConstraint):
+    def __init__(self):
+        super().__init__(-math.pi, math.pi)
+
+
+class _UnitInterval(_Interval, _SingletonConstraint):
+    def __init__(self):
+        super().__init__(0.0, 1.0)
+
+
 class _OpenInterval(_Interval):
     def __call__(self, x):
         return (x > self.lower_bound) & (x < self.upper_bound)
@@ -379,12 +427,7 @@ def __repr__(self):
         return fmt_string
 
 
-class _Circular(_Interval):
-    def __init__(self):
-        super().__init__(-math.pi, math.pi)
-
-
-class _LowerCholesky(Constraint):
+class _LowerCholesky(_SingletonConstraint):
     event_dim = 2
 
     def __call__(self, x):
@@ -420,7 +463,7 @@ def feasible_like(self, prototype):
         return jax.numpy.broadcast_to(value, prototype.shape)
 
 
-class _L1Ball(Constraint):
+class _L1Ball(_SingletonConstraint):
     """
     Constrain to the L1 ball of any dimension.
     """
@@ -437,7 +480,7 @@ def feasible_like(self, prototype):
         return jax.numpy.zeros_like(prototype)
 
 
-class _OrderedVector(Constraint):
+class _OrderedVector(_SingletonConstraint):
     event_dim = 1
 
     def __call__(self, x):
@@ -449,7 +492,7 @@ def feasible_like(self, prototype):
         )
 
 
-class _PositiveDefinite(Constraint):
+class _PositiveDefinite(_SingletonConstraint):
     event_dim = 2
 
     def __call__(self, x):
@@ -466,7 +509,7 @@ def feasible_like(self, prototype):
         )
 
 
-class _PositiveOrderedVector(Constraint):
+class _PositiveOrderedVector(_SingletonConstraint):
     """
     Constrains to a positive real-valued tensor where the elements are monotonically
     increasing along the `event_shape` dimension.
@@ -483,7 +526,7 @@ def feasible_like(self, prototype):
         )
 
 
-class _Real(Constraint):
+class _Real(_SingletonConstraint):
     def __call__(self, x):
         # XXX: consider to relax this condition to [-inf, inf] interval
         return (x == x) & (x != float("inf")) & (x != float("-inf"))
@@ -492,7 +535,7 @@ def feasible_like(self, prototype):
         return jax.numpy.zeros_like(prototype)
 
 
-class _Simplex(Constraint):
+class _Simplex(_SingletonConstraint):
     event_dim = 1
 
     def __call__(self, x):
@@ -503,7 +546,7 @@ def feasible_like(self, prototype):
         return jax.numpy.full_like(prototype, 1 / prototype.shape[-1])
 
 
-class _SoftplusPositive(_GreaterThan):
+class _SoftplusPositive(_GreaterThan, _SingletonConstraint):
     def __init__(self):
         super().__init__(lower_bound=0.0)
 
@@ -522,7 +565,7 @@ class _ScaledUnitLowerCholesky(_LowerCholesky):
     pass
 
 
-class _Sphere(Constraint):
+class _Sphere(_SingletonConstraint):
     """
     Constrain to the Euclidean sphere of any dimension.
     """
@@ -559,18 +602,18 @@ def feasible_like(self, prototype):
 lower_cholesky = _LowerCholesky()
 scaled_unit_lower_cholesky = _ScaledUnitLowerCholesky()
 multinomial = _Multinomial
-nonnegative_integer = _IntegerGreaterThan(0)
+nonnegative_integer = _IntegerNonnegative()
 ordered_vector = _OrderedVector()
-positive = _GreaterThan(0.0)
+positive = _Positive()
 positive_definite = _PositiveDefinite()
-positive_integer = _IntegerGreaterThan(1)
+positive_integer = _IntegerPositive()
 positive_ordered_vector = _PositiveOrderedVector()
 real = _Real()
-real_vector = independent(real, 1)
-real_matrix = independent(real, 2)
+real_vector = _RealVector()
+real_matrix = _RealMatrix()
 simplex = _Simplex()
 softplus_lower_cholesky = _SoftplusLowerCholesky()
 softplus_positive = _SoftplusPositive()
 sphere = _Sphere()
-unit_interval = _Interval(0.0, 1.0)
+unit_interval = _UnitInterval()
 open_interval = _OpenInterval
diff --git a/numpyro/distributions/transforms.py b/numpyro/distributions/transforms.py
--- a/numpyro/distributions/transforms.py
+++ b/numpyro/distributions/transforms.py
@@ -1065,6 +1065,11 @@ def _transform_to_corr_matrix(constraint):
     )
 
 
+@biject_to.register(type(constraints.positive))
+def _transform_to_positive(constraint):
+    return ExpTransform()
+
+
 @biject_to.register(constraints.greater_than)
 def _transform_to_greater_than(constraint):
     return ComposeTransform(
@@ -1085,6 +1090,8 @@ def _transform_to_less_than(constraint):
     )
 
 
+@biject_to.register(type(constraints.real_matrix))
+@biject_to.register(type(constraints.real_vector))
 @biject_to.register(constraints.independent)
 def _biject_to_independent(constraint):
     return IndependentTransform(
@@ -1092,6 +1099,12 @@ def _biject_to_independent(constraint):
     )
 
 
+@biject_to.register(type(constraints.unit_interval))
+def _transform_to_unit_interval(constraint):
+    return SigmoidTransform()
+
+
+@biject_to.register(type(constraints.circular))
 @biject_to.register(constraints.open_interval)
 @biject_to.register(constraints.interval)
 def _transform_to_interval(constraint):


Test patch:
diff --git a/test/test_pickle.py b/test/test_pickle.py
--- a/test/test_pickle.py
+++ b/test/test_pickle.py
@@ -13,6 +13,31 @@
 
 import numpyro
 import numpyro.distributions as dist
+from numpyro.distributions.constraints import (
+    boolean,
+    circular,
+    corr_cholesky,
+    corr_matrix,
+    greater_than,
+    interval,
+    l1_ball,
+    lower_cholesky,
+    nonnegative_integer,
+    ordered_vector,
+    positive,
+    positive_definite,
+    positive_integer,
+    positive_ordered_vector,
+    real,
+    real_matrix,
+    real_vector,
+    scaled_unit_lower_cholesky,
+    simplex,
+    softplus_lower_cholesky,
+    softplus_positive,
+    sphere,
+    unit_interval,
+)
 from numpyro.infer import (
     HMC,
     HMCECS,
@@ -93,3 +118,53 @@ def test_pickle_autoguide(guide_class):
     )
     samples = predictive(random.PRNGKey(1), None, 1)
     assert set(samples.keys()) == {"param", "x"}
+
+
+def test_pickle_singleton_constraint():
+    # some numpyro constraint classes such as constraints._Real, are only accessible
+    # through their public singleton instance, (such as constraint.real). This test
+    # ensures that pickling and unpickling singleton instances does not re-create
+    # additional instances, which is the default behavior of pickle, and which would
+    # break singleton semantics.
+    singleton_constraints = (
+        boolean,
+        circular,
+        corr_cholesky,
+        corr_matrix,
+        l1_ball,
+        lower_cholesky,
+        nonnegative_integer,
+        ordered_vector,
+        positive,
+        positive_definite,
+        positive_integer,
+        positive_ordered_vector,
+        real,
+        real_matrix,
+        real_vector,
+        scaled_unit_lower_cholesky,
+        simplex,
+        softplus_lower_cholesky,
+        softplus_positive,
+        sphere,
+        unit_interval,
+    )
+    for cnstr in singleton_constraints:
+        roundtripped_cnstr = pickle.loads(pickle.dumps(cnstr))
+        # make sure that the unpickled constraint is the original singleton constraint
+        assert roundtripped_cnstr is cnstr
+
+    # Test that it remains possible to pickle newly-created, non-singleton constraints.
+    # because these constraints are neither singleton nor exposed as top-level variables
+    # of the numpyro.distributions.constraints module, these objects are not pickled by
+    # reference, but by value.
+    int_cstr = interval(1.0, 2.0)
+    roundtripped_int_cstr = pickle.loads(pickle.dumps(int_cstr))
+    assert type(roundtripped_int_cstr) is type(int_cstr)
+    assert int_cstr.lower_bound == roundtripped_int_cstr.lower_bound
+    assert int_cstr.upper_bound == roundtripped_int_cstr.upper_bound
+
+    gt_cstr = greater_than(1.0)
+    roundtripped_gt_cstr = pickle.loads(pickle.dumps(gt_cstr))
+    assert type(roundtripped_gt_cstr) is type(gt_cstr)
+    assert gt_cstr.lower_bound == roundtripped_gt_cstr.lower_bound