Adopt singleton semantics for globally defined constraint instances (#1507)

* pickle constraints.real by reference

* test constraints.real pickling

* enforce singleton behavior for all singleton classes

* ensure pickling by reference for parametrized constraints

* remove singleton check against non-singleton constraint

* singleton subclasses for global parametrized constraints

* revert changes of #1519

* address old review comments

* separate out bijector optimizations

* minor fixes

* remove unused import

Author: pierreglaser

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

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,13 +1090,21 @@ 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(
         biject_to(constraint.base_constraint), constraint.reinterpreted_batch_ndims
     )
 
 
+@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/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