Repository: pyro-ppl/numpyro

Issue:
Some custom primitives does not work with the newest version of jax
The error happens because new version of jax requires a specific pattern for using `jit`'s static_argnums. See https://github.com/google/jax/pull/595 for more context. The failed functions include:

- [ ] xlogy
- [ ] xlog1py
- [ ] standard_gamma

In addition, I think that we can simplify the implementation of these functions by using decorator `custom_transform` as in `cumsum`, `cumprod`.

cc @neerajprad 


Discussion:
Are these functions failing with the jax master? We aren't using jit at least in `xlogy` and `xlog1py`, so that doesn't seem to immediately conflict with the JAX change. It also seems like a bunch of functions have been reorganized, so its probably worth tackling in one shot with the next jax release?
Yes, they fail with jax master. Maybe due to another reason. Yes, I also think that we just need to address these kinds of issues with one shot before our alpha release. It is almost there I guess. :)

Resolving patch:
diff --git a/numpyro/contrib/distributions/multivariate.py b/numpyro/contrib/distributions/multivariate.py
--- a/numpyro/contrib/distributions/multivariate.py
+++ b/numpyro/contrib/distributions/multivariate.py
@@ -15,7 +15,10 @@
 from numpyro.contrib.distributions.discrete import binom
 from numpyro.contrib.distributions.distribution import jax_continuous, jax_discrete, jax_multivariate
 from numpyro.distributions import constraints
-from numpyro.distributions.util import categorical_rvs, entr, multinomial_rvs, standard_gamma, xlogy
+from numpyro.distributions.util import categorical as categorical_rvs
+from numpyro.distributions.util import entr
+from numpyro.distributions.util import multinomial as multinomial_rvs
+from numpyro.distributions.util import standard_gamma, xlogy
 
 
 def _lnB(alpha):
@@ -148,7 +151,7 @@ def entropy(self, n, p):
     def _rvs(self, n, p):
         if self.is_logits:
             p = softmax(p)
-        return multinomial_rvs(self._random_state, n, p, self._size)
+        return multinomial_rvs(self._random_state, p, n, self._size)
 
 
 categorical = categorical_gen(name='categorical')
diff --git a/numpyro/distributions/discrete.py b/numpyro/distributions/discrete.py
--- a/numpyro/distributions/discrete.py
+++ b/numpyro/distributions/discrete.py
@@ -33,11 +33,11 @@
 from numpyro.distributions.util import (
     binary_cross_entropy_with_logits,
     binomial,
-    categorical_rvs,
+    categorical,
     clamp_probs,
     get_dtypes,
     lazy_property,
-    multinomial_rvs,
+    multinomial,
     poisson,
     promote_shapes,
     xlog1py,
@@ -215,7 +215,7 @@ def __init__(self, probs, total_count=1, validate_args=None):
                                           validate_args=validate_args)
 
     def sample(self, key, size=()):
-        return multinomial_rvs(key, self.total_count, self.probs, shape=size + self.batch_shape)
+        return multinomial(key, self.probs, self.total_count, shape=size + self.batch_shape)
 
     def log_prob(self, value):
         if self._validate_args:
@@ -256,7 +256,7 @@ def __init__(self, logits, total_count=1, validate_args=None):
                                                     validate_args=validate_args)
 
     def sample(self, key, size=()):
-        return multinomial_rvs(key, self.total_count, self.probs, shape=size + self.batch_shape)
+        return multinomial(key, self.probs, self.total_count, shape=size + self.batch_shape)
 
     def log_prob(self, value):
         if self._validate_args:
@@ -296,7 +296,7 @@ def __init__(self, probs, validate_args=None):
                                           validate_args=validate_args)
 
     def sample(self, key, size=()):
-        return categorical_rvs(key, self.probs, shape=size + self.batch_shape)
+        return categorical(key, self.probs, shape=size + self.batch_shape)
 
     def log_prob(self, value):
         if self._validate_args:
@@ -333,7 +333,7 @@ def __init__(self, logits, validate_args=None):
                                                     validate_args=validate_args)
 
     def sample(self, key, size=()):
-        return categorical_rvs(key, self.probs, shape=size + self.batch_shape)
+        return categorical(key, self.probs, shape=size + self.batch_shape)
 
     def log_prob(self, value):
         if self._validate_args:
diff --git a/numpyro/distributions/util.py b/numpyro/distributions/util.py
--- a/numpyro/distributions/util.py
+++ b/numpyro/distributions/util.py
@@ -7,9 +7,8 @@
 
 import jax.numpy as np
 from jax import canonicalize_dtype, custom_transforms, jit, lax, random, vmap
-from jax.core import Primitive
-from jax.interpreters import ad, partial_eval, xla
-from jax.numpy.lax_numpy import _promote_args_like, _promote_shapes
+from jax.interpreters import ad
+from jax.numpy.lax_numpy import _promote_args_like
 from jax.util import partial
 
 
@@ -168,159 +167,35 @@ def _standard_gamma_grad(sample, alpha):
     return grads.reshape(alpha.shape)
 
 
-def _standard_gamma_abstract_eval(key, alpha, jaxpr, aval, consts):
-    return lax.maybe_tracer_tuple_to_abstract_tuple(aval)
-
-
-def _standard_gamma_translate(c, key, alpha, jaxpr, aval, consts):
-    xla_computation = xla.jaxpr_computation(jaxpr, consts, (), c.GetShape(key), c.GetShape(alpha))
-    return c.Call(xla_computation, (key, alpha))
+@custom_transforms
+def _standard_gamma_p(key, alpha):
+    return _standard_gamma_impl(key, alpha)
 
 
-# define primitive
-standard_gamma_p = Primitive('standard_gamma')
-standard_gamma_p.def_impl(partial(xla.apply_primitive, standard_gamma_p))
-standard_gamma_p.def_abstract_eval(_standard_gamma_abstract_eval)
-xla.translations[standard_gamma_p] = _standard_gamma_translate
-ad.defjvp2(standard_gamma_p, None,
+ad.defjvp2(_standard_gamma_p.primitive, None,
            lambda tangent, sample, key, alpha, **kwargs: tangent * _standard_gamma_grad(sample, alpha))
 
 
 @partial(jit, static_argnums=(2, 3))
-def standard_gamma(key, alpha, shape=(), dtype=np.float32):
+def _standard_gamma(key, alpha, shape=(), dtype=np.float32):
     shape = shape or np.shape(alpha)
     alpha = lax.convert_element_type(alpha, dtype)
     if np.shape(alpha) != shape:
         alpha = np.broadcast_to(alpha, shape)
-    jaxpr, out, consts = partial_eval.trace_unwrapped_to_jaxpr(_standard_gamma_impl,
-                                                               tuple(lax._abstractify(o) for o in (key, alpha)))
-    aval, _ = out
-    return standard_gamma_p.bind(key, alpha, jaxpr=jaxpr, aval=aval, consts=consts)
-
-
-# TODO @npradhan: move to jax repo, if possible.
-def xlogy(x, y):
-    jaxpr, out, consts = partial_eval.trace_unwrapped_to_jaxpr(xlogy_impl, tuple(lax._abstractify(o) for o in (x, y)))
-    aval, _ = out
-    return xlogy_p.bind(x, y, jaxpr=jaxpr, aval=aval, consts=consts)
-
-
-def xlogy_impl(x, y):
-    return x * np.where(x == 0., 0., np.log(y))
-
-
-def xlogy_abstract_eval(x, y, jaxpr, aval, consts):
-    return lax.maybe_tracer_tuple_to_abstract_tuple(aval)
-
-
-def xlogy_translate(c, x, y, jaxpr, aval, consts):
-    xla_computation = xla.jaxpr_computation(jaxpr, consts, (), c.GetShape(x), c.GetShape(y))
-    return c.Call(xla_computation, (x, y))
-
-
-def xlogy_jvp_lhs(g, x, y, jaxpr, aval, consts):
-    x, y = _promote_args_like(osp_special.xlogy, x, y)
-    g, y = _promote_args_like(osp_special.xlogy, g, y)
-    return lax._safe_mul(lax._brcast(g, y), lax._brcast(lax.log(y), g))
-
-
-def xlogy_jvp_rhs(g, x, y, jaxpr, aval, consts):
-    x, y = _promote_args_like(osp_special.xlogy, x, y)
-    g, x = _promote_args_like(osp_special.xlogy, g, x)
-    jac = lax._safe_mul(lax._brcast(x, y), lax._brcast(lax.reciprocal(y), x))
-    return lax.mul(lax._brcast(g, jac), jac)
-
-
-xlogy_p = Primitive('xlogy')
-xlogy_p.def_impl(partial(xla.apply_primitive, xlogy_p))
-xlogy_p.def_abstract_eval(xlogy_abstract_eval)
-xla.translations[xlogy_p] = xlogy_translate
-ad.defjvp(xlogy_p, xlogy_jvp_lhs, xlogy_jvp_rhs)
-
-
-def xlog1py(x, y):
-    jaxpr, out, consts = partial_eval.trace_unwrapped_to_jaxpr(xlog1py_impl, tuple(lax._abstractify(o) for o in (x, y)))
-    aval, _ = out
-    return xlog1py_p.bind(x, y, jaxpr=jaxpr, aval=aval, consts=consts)
-
-
-def xlog1py_impl(x, y):
-    return x * np.where(x == 0., 0., np.log1p(y))
-
-
-def xlog1py_jvp_lhs(g, x, y, jaxpr, aval, consts):
-    x, y = _promote_args_like(osp_special.xlog1py, x, y)
-    g, y = _promote_args_like(osp_special.xlog1py, g, y)
-    return lax._safe_mul(lax._brcast(g, y), lax._brcast(lax.log1p(y), g))
-
-
-def xlog1py_jvp_rhs(g, x, y, jaxpr, aval, consts):
-    x, y = _promote_args_like(osp_special.xlog1py, x, y)
-    g, x = _promote_args_like(osp_special.xlog1py, g, x)
-    jac = lax._safe_mul(lax._brcast(x, y), lax._brcast(lax.reciprocal(1 + y), x))
-    return lax.mul(lax._brcast(g, jac), jac)
-
-
-xlog1py_p = Primitive('xlog1py')
-xlog1py_p.def_impl(partial(xla.apply_primitive, xlog1py_p))
-xlog1py_p.def_abstract_eval(xlogy_abstract_eval)
-xla.translations[xlog1py_p] = xlogy_translate
-ad.defjvp(xlog1py_p, xlog1py_jvp_lhs, xlog1py_jvp_rhs)
-
-
-def entr(p):
-    e = np.where(p > 0, -p * np.log(p), p)
-    e = np.where(e == 0, 0, e)
-    e = np.where(e < 0, -np.inf)
-    return e
-
+    return _standard_gamma_p(key, alpha)
 
-def binary_cross_entropy_with_logits(x, y):
-    # compute -y * log(sigmoid(x)) - (1 - y) * log(1 - sigmoid(x))
-    # Ref: https://www.tensorflow.org/api_docs/python/tf/nn/sigmoid_cross_entropy_with_logits
-    return np.clip(x, 0) + np.log1p(np.exp(-np.abs(x))) - x * y
-
-
-@custom_transforms
-def cumsum(x):
-    return np.cumsum(x, axis=-1)
-
-
-ad.defjvp(cumsum.primitive, lambda g, x: np.cumsum(g, axis=-1))
-
-
-@custom_transforms
-def cumprod(x):
-    return np.cumprod(x, axis=-1)
-
-
-# XXX this implementation does not address the case x=0, hence the result in that case will be nan
-# Ref: https://stackoverflow.com/questions/40916955/how-to-compute-gradient-of-cumprod-safely
-ad.defjvp2(cumprod.primitive, lambda g, ans, x: np.cumsum(g / x, axis=-1) * ans)
-
-
-def promote_shapes(*args, shape=()):
-    # adapted from lax.lax_numpy
-    if len(args) < 2 and not shape:
-        return args
-    else:
-        shapes = [np.shape(arg) for arg in args]
-        num_dims = len(lax.broadcast_shapes(shape, *shapes))
-        return [lax.reshape(arg, (1,) * (num_dims - len(s)) + s)
-                if len(s) < num_dims else arg for arg, s in zip(args, shapes)]
 
-
-def get_dtypes(*args):
-    return [canonicalize_dtype(onp.result_type(arg)) for arg in args]
+def standard_gamma(key, alpha, shape=(), dtype=np.float32):
+    return _standard_gamma(key, alpha, shape, dtype)
 
 
 # TODO: inefficient implementation; jit currently fails due to
 # dynamic size of random.uniform.
-@partial(jit, static_argnums=(0, 2, 3))
-def binomial(key, p, n=1, shape=()):
-    p, n = _promote_shapes(p, n)
+@partial(jit, static_argnums=(2, 3))
+def _binomial(key, p, n=1, shape=()):
+    p, n = promote_shapes(p, n)
     shape = shape or lax.broadcast_shapes(np.shape(p), np.shape(n))
-    n_max = np.max(n)
+    n_max = int(np.max(n))
     uniforms = random.uniform(key, shape + (n_max,))
     n = np.expand_dims(n, axis=-1)
     p = np.expand_dims(p, axis=-1)
@@ -329,8 +204,12 @@ def binomial(key, p, n=1, shape=()):
     return np.sum(mask * lax.lt(uniforms, p), axis=-1, keepdims=False)
 
 
+def binomial(key, p, n=1, shape=()):
+    return _binomial(key, p, n, shape)
+
+
 @partial(jit, static_argnums=(2,))
-def categorical_rvs(key, p, shape=()):
+def _categorical(key, p, shape=()):
     # this implementation is fast when event shape is small, and slow otherwise
     # Ref: https://stackoverflow.com/a/34190035
     shape = shape or p.shape[:-1]
@@ -341,8 +220,12 @@ def categorical_rvs(key, p, shape=()):
     return np.sum(s < r, axis=-1)
 
 
-@partial(jit, static_argnums=(1, 2))
-def poisson(key, rate, shape=()):
+def categorical(key, p, shape=()):
+    return _categorical(key, p, shape)
+
+
+@partial(jit, static_argnums=(2,))
+def _poisson(key, rate, shape=()):
     # Ref: https://en.wikipedia.org/wiki/Poisson_distribution#Generating_Poisson-distributed_random_variables
     shape = shape or np.shape(rate)
     L = np.exp(-rate)
@@ -361,6 +244,10 @@ def body_fn(val):
     return k - 1
 
 
+def poisson(key, rate, shape):
+    return _poisson(key, rate, shape)
+
+
 def _scatter_add_one(operand, indices, updates):
     return lax.scatter_add(operand, indices, updates,
                            lax.ScatterDimensionNumbers(update_window_dims=(),
@@ -368,8 +255,8 @@ def _scatter_add_one(operand, indices, updates):
                                                        scatter_dims_to_operand_dims=(0,)))
 
 
-@partial(jit, static_argnums=(1, 3))
-def multinomial_rvs(key, n, p, shape=()):
+@partial(jit, static_argnums=(2, 3))
+def _multinomial(key, p, n, shape=()):
     if np.shape(n) != np.shape(p)[:-1]:
         broadcast_shape = lax.broadcast_shapes(np.shape(n), np.shape(p)[:-1])
         n = np.broadcast_to(n, broadcast_shape)
@@ -377,7 +264,7 @@ def multinomial_rvs(key, n, p, shape=()):
     shape = shape or p.shape[:-1]
     n_max = int(np.max(n))
     # get indices from categorical distribution then gather the result
-    indices = categorical_rvs(key, p, (n_max,) + shape)
+    indices = categorical(key, p, (n_max,) + shape)
     # mask out values when counts is heterogeneous
     if not isinstance(n, Number):
         mask = promote_shapes(np.arange(n_max) < np.expand_dims(n, -1), shape=shape + (n_max,))[0]
@@ -395,6 +282,103 @@ def multinomial_rvs(key, n, p, shape=()):
     return np.reshape(samples_2D, shape + p.shape[-1:]) - excess
 
 
+def multinomial(key, p, n, shape=()):
+    return _multinomial(key, p, n, shape)
+
+
+def _xlogy_jvp_lhs(g, x, y):
+    shape = lax.broadcast_shapes(np.shape(g), np.shape(y))
+    g = np.broadcast_to(g, shape)
+    y = np.broadcast_to(y, shape)
+    g, y = _promote_args_like(osp_special.xlogy, g, y)
+    return lax._safe_mul(g, np.log(y))
+
+
+def _xlogy_jvp_rhs(g, x, y):
+    shape = lax.broadcast_shapes(np.shape(g), np.shape(x))
+    g = np.broadcast_to(g, shape)
+    x = np.broadcast_to(x, shape)
+    x, y = _promote_args_like(osp_special.xlogy, x, y)
+    return g * lax._safe_mul(x, np.reciprocal(y))
+
+
+@custom_transforms
+def xlogy(x, y):
+    x, y = _promote_args_like(osp_special.xlogy, x, y)
+    return lax._safe_mul(x, np.log(y))
+
+
+ad.defjvp(xlogy.primitive, _xlogy_jvp_lhs, _xlogy_jvp_rhs)
+
+
+def _xlog1py_jvp_lhs(g, x, y):
+    shape = lax.broadcast_shapes(np.shape(g), np.shape(y))
+    g = np.broadcast_to(g, shape)
+    y = np.broadcast_to(y, shape)
+    g, y = _promote_args_like(osp_special.xlog1py, g, y)
+    return lax._safe_mul(g, np.log1p(y))
+
+
+def _xlog1py_jvp_rhs(g, x, y):
+    shape = lax.broadcast_shapes(np.shape(g), np.shape(x))
+    g = np.broadcast_to(g, shape)
+    x = np.broadcast_to(x, shape)
+    x, y = _promote_args_like(osp_special.xlog1py, x, y)
+    return g * lax._safe_mul(x, np.reciprocal(1 + y))
+
+
+@custom_transforms
+def xlog1py(x, y):
+    x, y = _promote_args_like(osp_special.xlog1py, x, y)
+    return lax._safe_mul(x, np.log1p(y))
+
+
+ad.defjvp(xlog1py.primitive, _xlog1py_jvp_lhs, _xlog1py_jvp_rhs)
+
+
+def entr(p):
+    return np.where(p < 0, -np.inf, -xlogy(p))
+
+
+def binary_cross_entropy_with_logits(x, y):
+    # compute -y * log(sigmoid(x)) - (1 - y) * log(1 - sigmoid(x))
+    # Ref: https://www.tensorflow.org/api_docs/python/tf/nn/sigmoid_cross_entropy_with_logits
+    return np.clip(x, 0) + np.log1p(np.exp(-np.abs(x))) - x * y
+
+
+@custom_transforms
+def cumsum(x):
+    return np.cumsum(x, axis=-1)
+
+
+ad.defjvp(cumsum.primitive, lambda g, x: np.cumsum(g, axis=-1))
+
+
+@custom_transforms
+def cumprod(x):
+    return np.cumprod(x, axis=-1)
+
+
+# XXX this implementation does not address the case x=0, hence the result in that case will be nan
+# Ref: https://stackoverflow.com/questions/40916955/how-to-compute-gradient-of-cumprod-safely
+ad.defjvp2(cumprod.primitive, lambda g, ans, x: np.cumsum(g / x, axis=-1) * ans)
+
+
+def promote_shapes(*args, shape=()):
+    # adapted from lax.lax_numpy
+    if len(args) < 2 and not shape:
+        return args
+    else:
+        shapes = [np.shape(arg) for arg in args]
+        num_dims = len(lax.broadcast_shapes(shape, *shapes))
+        return [lax.reshape(arg, (1,) * (num_dims - len(s)) + s)
+                if len(s) < num_dims else arg for arg, s in zip(args, shapes)]
+
+
+def get_dtypes(*args):
+    return [canonicalize_dtype(onp.result_type(arg)) for arg in args]
+
+
 def sum_rightmost(x, dim):
     return np.sum(x, axis=-1) if dim == 1 else x
 


Test patch:
diff --git a/test/test_distributions.py b/test/test_distributions.py
--- a/test/test_distributions.py
+++ b/test/test_distributions.py
@@ -15,7 +15,7 @@
 import numpyro.distributions.constraints as constraints
 from numpyro.distributions.constraints import biject_to, matrix_to_tril_vec, vec_to_tril_matrix
 from numpyro.distributions.discrete import _to_probs_bernoulli, _to_probs_multinom
-from numpyro.distributions.util import multinomial_rvs, poisson
+from numpyro.distributions.util import multinomial, poisson
 
 
 def _identity(x): return x
@@ -137,7 +137,7 @@ def gen_values_within_bounds(constraint, size, key=random.PRNGKey(11)):
         return osp.dirichlet.rvs(alpha=np.ones((size[-1],)), size=size[:-1])
     elif isinstance(constraint, constraints._Multinomial):
         n = size[-1]
-        return multinomial_rvs(key, n=constraint.upper_bound, p=np.ones((n,)) / n, shape=size[:-1])
+        return multinomial(key, p=np.ones((n,)) / n, n=constraint.upper_bound, shape=size[:-1])
     else:
         raise NotImplementedError('{} not implemented.'.format(constraint))
 
@@ -161,7 +161,7 @@ def gen_values_outside_bounds(constraint, size, key=random.PRNGKey(11)):
         return osp.dirichlet.rvs(alpha=np.ones((size[-1],)), size=size[:-1]) + 1e-2
     elif isinstance(constraint, constraints._Multinomial):
         n = size[-1]
-        return multinomial_rvs(key, n=constraint.upper_bound, p=np.ones((n,)) / n, shape=size[:-1]) + 1
+        return multinomial(key, p=np.ones((n,)) / n, n=constraint.upper_bound, shape=size[:-1]) + 1
     else:
         raise NotImplementedError('{} not implemented.'.format(constraint))
 
@@ -242,6 +242,14 @@ def test_log_prob(jax_dist, sp_dist, params, prepend_shape, jit):
         expected = sp_dist.logpdf(samples)
     except AttributeError:
         expected = sp_dist.logpmf(samples)
+    except ValueError as e:
+        # precision issue: np.sum(x / np.sum(x)) = 0.99999994 != 1
+        if "The input vector 'x' must lie within the normal simplex." in str(e):
+            samples = samples.copy().astype('float64')
+            samples = samples / samples.sum(axis=-1, keepdims=True)
+            expected = sp_dist.logpdf(samples)
+        else:
+            raise e
     assert_allclose(jit_fn(jax_dist.log_prob)(samples), expected, atol=1e-5)
 
 
diff --git a/test/test_distributions_util.py b/test/test_distributions_util.py
--- a/test/test_distributions_util.py
+++ b/test/test_distributions_util.py
@@ -13,10 +13,10 @@
 
 from numpyro.distributions.util import (
     binary_cross_entropy_with_logits,
-    categorical_rvs,
+    categorical,
     cumprod,
     cumsum,
-    multinomial_rvs,
+    multinomial,
     standard_gamma,
     vec_to_tril_matrix,
     xlog1py,
@@ -157,7 +157,7 @@ def test_standard_gamma_grad(alpha):
     pdf = osp_stats.gamma.pdf(z, alpha)
     expected_grad = -cdf_dot / pdf
 
-    assert_allclose(actual_grad, expected_grad, rtol=0.0005)
+    assert_allclose(actual_grad, expected_grad, atol=1e-8, rtol=0.0005)
 
 
 @pytest.mark.parametrize('p, shape', [
@@ -169,7 +169,7 @@ def test_standard_gamma_grad(alpha):
 def test_categorical_shape(p, shape):
     rng = random.PRNGKey(0)
     expected_shape = lax.broadcast_shapes(p.shape[:-1], shape)
-    assert np.shape(categorical_rvs(rng, p, shape)) == expected_shape
+    assert np.shape(categorical(rng, p, shape)) == expected_shape
 
 
 @pytest.mark.parametrize("p", [
@@ -179,7 +179,7 @@ def test_categorical_shape(p, shape):
 def test_categorical_stats(p):
     rng = random.PRNGKey(0)
     n = 10000
-    z = categorical_rvs(rng, p, (n,))
+    z = categorical(rng, p, (n,))
     _, counts = onp.unique(z, return_counts=True)
     assert_allclose(counts / float(n), p, atol=0.01)
 
@@ -194,7 +194,7 @@ def test_multinomial_shape(p, shape):
     rng = random.PRNGKey(0)
     n = 10000
     expected_shape = lax.broadcast_shapes(p.shape[:-1], shape) + p.shape[-1:]
-    assert np.shape(multinomial_rvs(rng, n, p, shape)) == expected_shape
+    assert np.shape(multinomial(rng, p, n, shape)) == expected_shape
 
 
 @pytest.mark.parametrize("p", [
@@ -207,7 +207,7 @@ def test_multinomial_shape(p, shape):
 ])
 def test_multinomial_stats(p, n):
     rng = random.PRNGKey(0)
-    z = multinomial_rvs(rng, n, p)
+    z = multinomial(rng, p, n)
     n = float(n) if isinstance(n, Number) else np.expand_dims(n.astype(p.dtype), -1)
     p = np.broadcast_to(p, z.shape)
     assert_allclose(z / n, p, atol=0.01)