Repository: pyro-ppl/numpyro

Issue:
Gradient propagation failure for masked invalid values
I am working on a model that requires implementing a plate over arrays with variable sizes. My understanding is that the conventional solution to this problem is padding shorter arrays and then masking the padded values. 

I initially implemented this with `nan` pads, but found that that yielded completely divergent samples. 

With @fehiepsi's [help](https://forum.pyro.ai/t/hierarchical-bayesian-model-with-data-of-varying-length/1588), I was able to solve this problem by padding with a valid value instead of `nan`. They suggested that even though the padding values are masked, the gradient is still propagated incorrectly. 

I can confirm that parameter estimates seem to be invariant across masked valid values; however, masked invalid values (e.g. including `nan` and values excluded from the support of the relevant probability distribution) cause 100% divergence.

@fehiepsi asked me to open an issue and suggested a fix that automatically replaces invalid values with valid values.  

I included a reproducible example below. 

```


# Generate Data

# specify number of observations for each group
n1 = 2000
n2 = 1000

# Specify pad value
pad = -1

# Sample observations from different beta distributions for each group, given specified sample size
np.random.seed(seed=123)
y1 = beta.rvs(.5, .5, size=n1)
np.random.seed(seed=321)
y2 = beta.rvs(3, 1, size=n2)

pads = numpy.ones((n1-n2)) * pad
y2 = np.concatenate([y2, pads])

# Combine samples
y = np.array([y1, y2], order='K').T

sample_sizes =  np.array([n1, n2])

def model(sample_sizes, y = None):
    
    # Define the number of groups from which we've attained samples
    n_groups =  sample_sizes.shape[0]
    
    # Initialize a plate for each group
    with numpyro.plate("n_groups", n_groups):
        
        # Set priors on a and b within each group
        a = numpyro.sample('a', dist.Gamma(.5, .5))
        b =  numpyro.sample('b', dist.Gamma(3, 3))
        
        # Calculate maximum group size
        I = sample_sizes.max().item()
        # Create a range of 0:I
        i = torch.arange(I).unsqueeze(-1).numpy()
        
        # Initialize a plate for each observation
        with numpyro.plate('data', I):
            # Mask observations that exceed sample size in the respective plate
            with numpyro.handlers.mask(mask = i < sample_sizes): 
                # Estimate y_hat
                y_hat = numpyro.sample('y_hat', dist.Beta(a, b), obs=y)

rng_key = random.PRNGKey(0)
rng_key, rng_key_ = random.split(rng_key)

num_warmup, num_samples = 500, 2000

# Run NUTS.
kernel = NUTS(model)
mcmc = MCMC(kernel, num_warmup, num_samples, num_chains = 3)

start = time.time()
mcmc.run(rng_key_, sample_sizes=sample_sizes, y = y)
end = time.time()
mcmc.print_summary()
```


Discussion:

Resolving patch:
diff --git a/numpyro/distributions/constraints.py b/numpyro/distributions/constraints.py
--- a/numpyro/distributions/constraints.py
+++ b/numpyro/distributions/constraints.py
@@ -73,11 +73,20 @@ def check(self, value):
         """
         return self(value)
 
+    def feasible_like(self, prototype):
+        """
+        Get a feasible value which has the same shape as dtype as `prototype`.
+        """
+        raise NotImplementedError
+
 
 class _Boolean(Constraint):
     def __call__(self, x):
         return (x == 0) | (x == 1)
 
+    def feasible_like(self, prototype):
+        return jax.numpy.zeros_like(prototype)
+
 
 class _CorrCholesky(Constraint):
     def __call__(self, x):
@@ -89,6 +98,9 @@ def __call__(self, x):
         unit_norm_row = jnp.all((x_norm <= 1) & (x_norm > 1 - 1e-6), axis=-1)
         return lower_triangular & positive_diagonal & unit_norm_row
 
+    def feasible_like(self, prototype):
+        return jax.numpy.broadcast_to(jax.numpy.eye(prototype.shape[-1]), prototype.shape)
+
 
 class _CorrMatrix(Constraint):
     def __call__(self, x):
@@ -101,6 +113,9 @@ def __call__(self, x):
         unit_variance = jnp.all(jnp.abs(jnp.diagonal(x, axis1=-2, axis2=-1) - 1) < 1e-6, axis=-1)
         return symmetric & positive & unit_variance
 
+    def feasible_like(self, prototype):
+        return jax.numpy.broadcast_to(jax.numpy.eye(prototype.shape[-1]), prototype.shape)
+
 
 class _Dependent(Constraint):
     def __call__(self, x):
@@ -118,6 +133,9 @@ def __init__(self, lower_bound):
     def __call__(self, x):
         return x > self.lower_bound
 
+    def feasible_like(self, prototype):
+        return jax.numpy.broadcast_to(self.lower_bound + 1, jax.numpy.shape(prototype))
+
 
 class _LessThan(Constraint):
     def __init__(self, upper_bound):
@@ -126,6 +144,9 @@ def __init__(self, upper_bound):
     def __call__(self, x):
         return x < self.upper_bound
 
+    def feasible_like(self, prototype):
+        return jax.numpy.broadcast_to(self.upper_bound - 1, jax.numpy.shape(prototype))
+
 
 class _IntegerInterval(Constraint):
     def __init__(self, lower_bound, upper_bound):
@@ -135,6 +156,9 @@ def __init__(self, lower_bound, upper_bound):
     def __call__(self, x):
         return (x >= self.lower_bound) & (x <= self.upper_bound) & (x % 1 == 0)
 
+    def feasible_like(self, prototype):
+        return jax.numpy.broadcast_to(self.lower_bound, jax.numpy.shape(prototype))
+
 
 class _IntegerGreaterThan(Constraint):
     def __init__(self, lower_bound):
@@ -143,6 +167,9 @@ def __init__(self, lower_bound):
     def __call__(self, x):
         return (x % 1 == 0) & (x >= self.lower_bound)
 
+    def feasible_like(self, prototype):
+        return jax.numpy.broadcast_to(self.lower_bound, jax.numpy.shape(prototype))
+
 
 class _Interval(Constraint):
     def __init__(self, lower_bound, upper_bound):
@@ -152,6 +179,9 @@ def __init__(self, lower_bound, upper_bound):
     def __call__(self, x):
         return (x >= self.lower_bound) & (x <= self.upper_bound)
 
+    def feasible_like(self, prototype):
+        return jax.numpy.broadcast_to((self.lower_bound + self.upper_bound) / 2, jax.numpy.shape(prototype))
+
 
 class _LowerCholesky(Constraint):
     def __call__(self, x):
@@ -161,6 +191,9 @@ def __call__(self, x):
         positive_diagonal = jnp.all(jnp.diagonal(x, axis1=-2, axis2=-1) > 0, axis=-1)
         return lower_triangular & positive_diagonal
 
+    def feasible_like(self, prototype):
+        return jax.numpy.broadcast_to(jax.numpy.eye(prototype.shape[-1]), prototype.shape)
+
 
 class _Multinomial(Constraint):
     def __init__(self, upper_bound):
@@ -169,11 +202,19 @@ def __init__(self, upper_bound):
     def __call__(self, x):
         return (x >= 0).all(axis=-1) & (x.sum(axis=-1) == self.upper_bound)
 
+    def feasible_like(self, prototype):
+        pad_width = ((0, 0),) * jax.numpy.ndim(self.upper_bound) + ((0, prototype.shape[-1] - 1),)
+        value = jax.numpy.pad(jax.numpy.expand_dims(self.upper_bound, -1), pad_width)
+        return jax.numpy.broadcast_to(value, prototype.shape)
+
 
 class _OrderedVector(Constraint):
     def __call__(self, x):
         return (x[..., 1:] > x[..., :-1]).all(axis=-1)
 
+    def feasible_like(self, prototype):
+        return jax.numpy.broadcast_to(jax.numpy.arange(float(prototype.shape[-1])), prototype.shape)
+
 
 class _PositiveDefinite(Constraint):
     def __call__(self, x):
@@ -184,23 +225,35 @@ def __call__(self, x):
         positive = jnp.linalg.eigh(x)[0][..., 0] > 0
         return symmetric & positive
 
+    def feasible_like(self, prototype):
+        return jax.numpy.broadcast_to(jax.numpy.eye(prototype.shape[-1]), prototype.shape)
+
 
 class _Real(Constraint):
     def __call__(self, x):
         # XXX: consider to relax this condition to [-inf, inf] interval
         return (x == x) & (x != float('inf')) & (x != float('-inf'))
 
+    def feasible_like(self, prototype):
+        return jax.numpy.zeros_like(prototype)
+
 
 class _RealVector(Constraint):
     def __call__(self, x):
         return ((x == x) & (x != float('inf')) & (x != float('-inf'))).all(axis=-1)
 
+    def feasible_like(self, prototype):
+        return jax.numpy.zeros_like(prototype)
+
 
 class _Simplex(Constraint):
     def __call__(self, x):
         x_sum = x.sum(axis=-1)
         return (x >= 0).all(axis=-1) & (x_sum < 1 + 1e-6) & (x_sum > 1 - 1e-6)
 
+    def feasible_like(self, prototype):
+        return jax.numpy.full_like(prototype, 1 / prototype.shape[-1])
+
 
 # TODO: Make types consistent
 
diff --git a/numpyro/distributions/distribution.py b/numpyro/distributions/distribution.py
--- a/numpyro/distributions/distribution.py
+++ b/numpyro/distributions/distribution.py
@@ -296,7 +296,7 @@ def to_event(self, reinterpreted_batch_ndims=None):
         """
         if reinterpreted_batch_ndims is None:
             reinterpreted_batch_ndims = len(self.batch_shape)
-        elif reinterpreted_batch_ndims == 0:
+        if reinterpreted_batch_ndims == 0:
             return self
         return Independent(self, reinterpreted_batch_ndims)
 
@@ -687,6 +687,13 @@ def log_prob(self, value):
             return jnp.zeros(shape)
         if self._mask is True:
             return self.base_dist.log_prob(value)
+        try:
+            default_value = self.base_dist.support.feasible_like(value)
+        except NotImplementedError:
+            pass
+        else:
+            mask = jnp.reshape(self._mask, jnp.shape(self._mask) + (1,) * self.event_dim)
+            value = jnp.where(mask, value, default_value)
         return jnp.where(self._mask, self.base_dist.log_prob(value), 0.)
 
     def enumerate_support(self, expand=True):


Test patch:
diff --git a/test/test_distributions.py b/test/test_distributions.py
--- a/test/test_distributions.py
+++ b/test/test_distributions.py
@@ -910,6 +910,17 @@ def g(x):
 def test_constraints(constraint, x, expected):
     assert_array_equal(constraint(x), expected)
 
+    feasible_value = constraint.feasible_like(x)
+    assert jnp.shape(feasible_value) == jnp.shape(x)
+    assert_allclose(constraint(feasible_value), jnp.full(jnp.shape(expected), True))
+
+    try:
+        inverse = biject_to(constraint).inv(feasible_value)
+    except NotImplementedError:
+        pass
+    else:
+        assert_allclose(inverse, jnp.zeros_like(inverse), atol=2e-7)
+
 
 @pytest.mark.parametrize('constraint', [
     constraints.corr_cholesky,
@@ -1302,6 +1313,20 @@ def test_mask(batch_shape, event_shape, mask_shape):
     assert_allclose(actual != 0, jnp.broadcast_to(mask, lax.broadcast_shapes(batch_shape, mask_shape)))
 
 
+@pytest.mark.parametrize('event_shape', [(), (4,), (2, 4)])
+def test_mask_grad(event_shape):
+    def f(x, data):
+        base_dist = dist.Beta(jnp.exp(x), jnp.ones(event_shape)).to_event()
+        mask = jnp.all(jnp.isfinite(data), tuple(-i-1 for i in range(len(event_shape))))
+        log_prob = base_dist.mask(mask).log_prob(data)
+        assert log_prob.shape == data.shape[:len(data.shape) - len(event_shape)]
+        return log_prob.sum()
+
+    data = jnp.array([[0.4, jnp.nan, 0.2, jnp.nan], [0.5, 0.5, 0.5, 0.5]])
+    log_prob, grad = jax.value_and_grad(f)(1., data)
+    assert jnp.isfinite(grad) and jnp.isfinite(log_prob)
+
+
 @pytest.mark.parametrize('jax_dist, sp_dist, params', CONTINUOUS + DISCRETE + DIRECTIONAL)
 def test_dist_pytree(jax_dist, sp_dist, params):
     def f(x):