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.

Changed files:
tests/scipy_stats_test.py, jax/_src/scipy/stats/dirichlet.py

diff --git a/jax/_src/scipy/stats/dirichlet.py b/jax/_src/scipy/stats/dirichlet.py
index f58aaf26b..2b8923708 100644
--- 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)
diff --git a/tests/scipy_stats_test.py b/tests/scipy_stats_test.py
index 7c3ce5b49..1635da564 100644
--- a/tests/scipy_stats_test.py
+++ b/tests/scipy_stats_test.py
@@ -30,6 +30,7 @@ from jax.config import config
 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 @@ class LaxBackedScipyStatsTests(jtu.JaxTestCase):
                             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):