Repository: google/jax

Issue:
JAX Dirichlet Log PDF output different thatn SciPy
```
jax.scipy.stats.dirichlet.logpdf(np.array([9.2820191190831968786056904718861915171146e-01,
1.7047643862547270932061849180172430351377e-03,
3.8956840931979931121065252597190919914283e-03,
5.3582474241543359694261994263797532767057e-02,
6.7548574969396027109502789187445159768686e-03,
2.8701918527402454049679558778507271199487e-03,
2.9901160210044600180900875585621179197915e-03]), np.array([2.0922036170959472656250000000000000000000e+00,
2.5644928216934204101562500000000000000000e-01,
1.0264828205108642578125000000000000000000e+00,
6.1082868576049804687500000000000000000000e+00,
5.9878396987915039062500000000000000000000e-01,
2.6282947063446044921875000000000000000000e+00,
7.7161555290222167968750000000000000000000e+00], dtype=np.float32))
```
outputs `-inf` whereas 
```
scipy.stats.dirichlet.logpdf(np.array([9.2820191190831968786056904718861915171146e-01,
1.7047643862547270932061849180172430351377e-03,
3.8956840931979931121065252597190919914283e-03,
5.3582474241543359694261994263797532767057e-02,
6.7548574969396027109502789187445159768686e-03,
2.8701918527402454049679558778507271199487e-03,
2.9901160210044600180900875585621179197915e-03]), np.array([2.0922036170959472656250000000000000000000e+00,
2.5644928216934204101562500000000000000000e-01,
1.0264828205108642578125000000000000000000e+00,
6.1082868576049804687500000000000000000000e+00,
5.9878396987915039062500000000000000000000e-01,
2.6282947063446044921875000000000000000000e+00,
7.7161555290222167968750000000000000000000e+00], dtype=np.float32))
```
outputs -31.38.

I tried using 64-bit floats but that didn't work either.


Discussion:
What version of JAX are you using? On version 0.2.9 I get the correct output.
Same version. It might be a precision issue. I can repro with:
```
from jax.config import config
config.update("jax_enable_x64", True)
```
at the top of the script and cannot without those additional lines.
Digging around the JAX code, I think it is this (line 38, `dirichlet.py`):
```
    return jnp.where(_is_simplex(x), log_probs, -jnp.inf)
```
where
```
def _is_simplex(x):
    x_sum = jnp.sum(x, axis=-1)
    return jnp.all(x > 0, axis=-1) & (x_sum <= 1) & (x_sum > 1 - 1e-6)
```
Something about the precision makes the sum of those fractions marginally greater than 1. I guess this counts as user error but I will say getting negative infinite is much worse than an exception (I thought my training was inserting non positive `alpha` parameters. 
Ah, I see. It looks like JAX and scipy use different criteria to check if the value is close enough to 1. Jax is here: https://github.com/google/jax/blob/ab21409e0b2a561f217c833ca6f915b9dcefffd0/jax/_src/scipy/stats/dirichlet.py#L27

And scipy is here: https://github.com/scipy/scipy/blob/fa119c98b18e3dfc6e158e22e806831f40d660d3/scipy/stats/_multivariate.py#L1287

Basically, jax defines the valid input range as `1 - 1E-6 < x_sum <= 1` and scipy defines it as `1 - 10E-10 <= x_sum <= 1 + 10E-10`

The fix here would be to make these constraints match.
Nice. 

The SciPy tolerance seems more forgiving, normalizing within the current JAX tolerance is trickier than it seems (dividing by the sum of the array isn't sufficient, so some form of truncating precision and then summing might have to do).
I started working on this, and it looks like the problems run even deeper. Due to a copy/paste typo, we never even test dirichlet distributions (note `cauchy` here): https://github.com/google/jax/blob/3c87a36831109ddcbc95e3844108c78f51d00300/tests/scipy_stats_test.py#L166-L170

And because of this we never discovered that the implementation in JAX is incompatible with the implementation in scipy in ways that are deeper than the tolerance around the sum of x.
For example:
```python
import numpy as np
import scipy.stats
import jax.scipy.stats

np.random.seed(0)
x = np.random.rand(4)
x /= x.sum()
alpha = np.ones(5)

print(scipy.stats.dirichlet.logpdf(x, alpha))
# 3.1780538303479458

print(jax.scipy.stats.dirichlet.logpdf(x, alpha))
# ValueError: Incompatible shapes for broadcasting: ((5,), (4,), (1,))
```

Resolving patch:
diff --git a/jax/_src/scipy/stats/dirichlet.py b/jax/_src/scipy/stats/dirichlet.py
--- a/jax/_src/scipy/stats/dirichlet.py
+++ b/jax/_src/scipy/stats/dirichlet.py
@@ -13,29 +13,40 @@
 # limitations under the License.
 
 
-import numpy as np
 import scipy.stats as osp_stats
 
 from jax import lax
 from jax._src.numpy import lax_numpy as jnp
 from jax._src.numpy.util import _wraps
 from jax.scipy.special import gammaln, xlogy
+from jax._src.numpy.lax_numpy import _promote_dtypes_inexact
 
 
 def _is_simplex(x):
-    x_sum = jnp.sum(x, axis=-1)
-    return jnp.all(x > 0, axis=-1) & (x_sum <= 1) & (x_sum > 1 - 1e-6)
+  x_sum = jnp.sum(x, axis=0)
+  return jnp.all(x > 0, axis=0) & (abs(x_sum - 1) < 1E-6)
 
 
 @_wraps(osp_stats.dirichlet.logpdf, update_doc=False)
 def logpdf(x, alpha):
-    args = (np.ones((0,), lax.dtype(x)), np.ones((1,), lax.dtype(alpha)))
-    to_dtype = lax.dtype(osp_stats.dirichlet.logpdf(*args))
-    x, alpha = [lax.convert_element_type(arg, to_dtype) for arg in (x, alpha)]
-    one = jnp._constant_like(x, 1)
-    normalize_term = jnp.sum(gammaln(alpha), axis=-1) - gammaln(jnp.sum(alpha, axis=-1))
-    log_probs = lax.sub(jnp.sum(xlogy(lax.sub(alpha, one), x), axis=-1), normalize_term)
-    return jnp.where(_is_simplex(x), log_probs, -jnp.inf)
+  x, alpha = _promote_dtypes_inexact(x, alpha)
+  if alpha.ndim != 1:
+    raise ValueError(
+      f"`alpha` must be one-dimensional; got alpha.shape={alpha.shape}"
+    )
+  if x.shape[0] not in (alpha.shape[0], alpha.shape[0] - 1):
+    raise ValueError(
+      "`x` must have either the same number of entries as `alpha` "
+      f"or one entry fewer; got x.shape={x.shape}, alpha.shape={alpha.shape}"
+    )
+  one = jnp._constant_like(x, 1)
+  if x.shape[0] != alpha.shape[0]:
+    x = jnp.concatenate([x, lax.sub(one, x.sum(0, keepdims=True))], axis=0)
+  normalize_term = jnp.sum(gammaln(alpha)) - gammaln(jnp.sum(alpha))
+  if x.ndim > 1:
+    alpha = lax.broadcast_in_dim(alpha, alpha.shape + (1,) * (x.ndim - 1), (0,))
+  log_probs = lax.sub(jnp.sum(xlogy(lax.sub(alpha, one), x), axis=0), normalize_term)
+  return jnp.where(_is_simplex(x), log_probs, -jnp.inf)
 
 
 @_wraps(osp_stats.dirichlet.pdf, update_doc=False)


Test patch:
diff --git a/tests/scipy_stats_test.py b/tests/scipy_stats_test.py
--- a/tests/scipy_stats_test.py
+++ b/tests/scipy_stats_test.py
@@ -30,6 +30,7 @@
 config.parse_flags_with_absl()
 
 all_shapes = [(), (4,), (3, 4), (3, 1), (1, 4), (2, 1, 4)]
+one_and_two_dim_shapes = [(4,), (3, 4), (3, 1), (1, 4)]
 scipy_version = tuple(map(int, osp.version.version.split('.')[:2]))
 
 
@@ -163,22 +164,44 @@ def args_maker():
                             tol=1e-4)
     self._CompileAndCheck(lax_fun, args_maker)
 
-  @genNamedParametersNArgs(2)
+  @parameterized.named_parameters(
+    jtu.cases_from_list(
+      {"testcase_name": jtu.format_test_name_suffix("", [x_shape, alpha_shape], dtypes),
+        "shapes": [x_shape, alpha_shape], "dtypes": dtypes}
+      for x_shape in one_and_two_dim_shapes
+      for alpha_shape in [(x_shape[0],), (x_shape[0] + 1,)]
+      for dtypes in itertools.combinations_with_replacement(jtu.dtypes.floating, 2)
+  ))
   def testDirichletLogPdf(self, shapes, dtypes):
     rng = jtu.rand_positive(self.rng())
-    scipy_fun = osp_stats.cauchy.logpdf
-    lax_fun = lsp_stats.cauchy.logpdf
-    dim = 4
-    shapes = (shapes[0] + (dim,), shapes[1] + (dim,))
+
+    def _normalize(x, alpha):
+      x_norm = x.sum(0) + (0.0 if x.shape[0] == alpha.shape[0] else 0.1)
+      return (x / x_norm).astype(x.dtype), alpha
+
+    def lax_fun(x, alpha):
+      return lsp_stats.dirichlet.logpdf(*_normalize(x, alpha))
+
+    def scipy_fun(x, alpha):
+      # scipy validates the x normalization using float64 arithmetic, so we must
+      # cast x to float64 before normalization to ensure this passes.
+      x, alpha = _normalize(x.astype('float64'), alpha)
+
+      result = osp_stats.dirichlet.logpdf(x, alpha)
+      # if x.shape is (N, 1), scipy flattens the output, while JAX returns arrays
+      # of a consistent rank. This check ensures the results have the same shape.
+      return result if x.ndim == 1 else np.atleast_1d(result)
 
     def args_maker():
+      # Don't normalize here, because we want normalization to happen at 64-bit
+      # precision in the scipy version.
       x, alpha = map(rng, shapes, dtypes)
-      x = x / np.sum(x, axis=-1, keepdims=True)
-      return [x, alpha]
+      return x, alpha
 
+    tol = {np.float32: 1E-3, np.float64: 1e-5}
     self._CheckAgainstNumpy(scipy_fun, lax_fun, args_maker, check_dtypes=False,
-                            tol=1e-4)
-    self._CompileAndCheck(lax_fun, args_maker)
+                            tol=tol)
+    self._CompileAndCheck(lax_fun, args_maker, atol=tol, rtol=tol)
 
   @genNamedParametersNArgs(3)
   def testExponLogPdf(self, shapes, dtypes):