Repository: pyro-ppl/numpyro

Issue:
a bug in SimplexToOrderedTransform
I run the tutorial Ordinal Regression (https://num.pyro.ai/en/latest/tutorials/ordinal_regression.html).
I am not sure how the transforms.SimplexToOrderedTransform behaves.

I have confirmed that it usually works as follows.

```
d = dist.TransformedDistribution(dist.Dirichlet(np.ones((3,))),dist.transforms.SimplexToOrderedTransform(0))
d.sample(random.PRNGKey(0))
```
```
Array([-1.584826 ,  0.6460422], dtype=float32)
```

However, when executing MCMC, it does not seem to work correctly as shown in the following code.
This is a bug in SimplexToOrderedTransform where the methods forward_shape and inverse_shape are not implemented correctly (it is currently using the default ones, that maps shape to shape rather than shape to shape[:-1] + (shape[-1] - 1,). 

```
import os

import jax.numpy as jnp
from jax import random
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns

import numpyro
from numpyro.diagnostics import hpdi
import numpyro.distributions as dist
from jax.experimental.ode import odeint
import numpyro.distributions.constraints as constraints
from numpyro.infer import MCMC, NUTS
from numpyro.infer import Predictive
from numpyro.infer.util import initialize_model

import arviz as az

az.style.use("arviz-darkgrid")

assert numpyro.__version__.startswith("0.11.0")

numpyro.set_platform("cpu")
numpyro.set_host_device_count(1)
from numpyro.infer.reparam import TransformReparam


# data generation
simkeys = random.split(random.PRNGKey(1), 2)
nsim = 50
nclasses = 3
Y = dist.Categorical(logits=np.zeros(nclasses)).sample(simkeys[0], sample_shape=(nsim,))
X =dist. Normal().sample(simkeys[1], sample_shape=(nsim,))
X += Y
df = pd.DataFrame({"X": X, "Y": Y})


def model_ng(X, Y, nclasses, concentration, anchor_point=0.0):
    b_X_eta = numpyro.sample("b_X_eta", dist.Normal(0, 5))

    #with numpyro.handlers.reparam(config={"c_y": TransformReparam()}):
    c_y = numpyro.sample(
        "c_y",
        dist.TransformedDistribution(
            dist.Dirichlet(concentration),
            dist.transforms.SimplexToOrderedTransform(anchor_point),
        )
    )
    print(c_y.shape, c_y)
    with numpyro.plate("obs", X.shape[0]):
        eta = X * b_X_eta
        numpyro.sample("Y", dist.OrderedLogistic(eta, c_y), obs=Y)


concentration = np.ones((nclasses,)) * 10.0

rng_key= random.PRNGKey(0)
kernel = NUTS(model_ng)
mcmc = MCMC(kernel, num_warmup=1000, num_samples=2000, num_chains=1, thinning=1)
mcmc.run(
    rng_key=rng_key,
    X=df["X"].values,
    Y=df["Y"].values,
    nclasses=nclasses,
    concentration=concentration,
)
# with exclude_deterministic=False, we will also show the ordinal probabilities sampled from Dirichlet (vis. `c_y_base`)
mcmc.print_summary(exclude_deterministic=False)
```
```
(3,) [1.2993407 2.6367652 9.292245 ]
(3,) Traced<ConcreteArray([1.2993407 2.6367652 9.292245 ], dtype=float32)>with<JVPTrace(level=2/0)> with
  primal = Array([1.2993407, 2.6367652, 9.292245 ], dtype=float32)
  tangent = Traced<ShapedArray(float32[3])>with<JaxprTrace(level=1/0)> with
    pval = (ShapedArray(float32[3]), None)
    recipe = JaxprEqnRecipe(eqn_id=<object object at 0x3019bfba0>, in_tracers=(Traced<ShapedArray(float32[3]):JaxprTrace(level=1/0)>,), out_tracer_refs=[<weakref at 0x3028781d0; to 'JaxprTracer' at 0x302878400>], out_avals=[ShapedArray(float32[3])], primitive=pjit, params={'jaxpr': { lambda ; a:f32[3]. let b:f32[3] = cumsum[axis=0 reverse=False] a in (b,) }, 'in_shardings': (UnspecifiedValue,), 'out_shardings': (UnspecifiedValue,), 'resource_env': None, 'donated_invars': (False,), 'name': '_cumulative_reduction', 'in_positional_semantics': (<_PositionalSemantics.GLOBAL: 1>,), 'out_positional_semantics': <_PositionalSemantics.GLOBAL: 1>, 'keep_unused': False, 'inline': False}, effects=set(), source_info=SourceInfo(traceback=<jaxlib.xla_extension.Traceback object at 0x174ac8330>, name_stack=NameStack(stack=(Transform(name='jvp'),))))
Output exceeds the size limit. Open the full output data in a text editor
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[58], line 62
     60 kernel = NUTS(model_ng)
     61 mcmc = MCMC(kernel, num_warmup=1000, num_samples=2000, num_chains=1, thinning=1)
---> 62 mcmc.run(
     63     rng_key=rng_key,
     64     X=df["X"].values,
     65     Y=df["Y"].values,
     66     nclasses=nclasses,
     67     concentration=concentration,
     68 )
     69 # with exclude_deterministic=False, we will also show the ordinal probabilities sampled from Dirichlet (vis. `c_y_base`)
     70 mcmc.print_summary(exclude_deterministic=False)

File ~/Desktop/programming/numpyro_intro/.venv/lib/python3.9/site-packages/numpyro/infer/mcmc.py:628, in MCMC.run(self, rng_key, extra_fields, init_params, *args, **kwargs)
    626 map_args = (rng_key, init_state, init_params)
    627 if self.num_chains == 1:
--> 628     states_flat, last_state = partial_map_fn(map_args)
    629     states = tree_map(lambda x: x[jnp.newaxis, ...], states_flat)
    630 else:

File ~/Desktop/programming/numpyro_intro/.venv/lib/python3.9/site-packages/numpyro/infer/mcmc.py:410, in MCMC._single_chain_mcmc(self, init, args, kwargs, collect_fields)
    408 # Check if _sample_fn is None, then we need to initialize the sampler.
    409 if init_state is None or (getattr(self.sampler, "_sample_fn", None) is None):
...
-> 1617       raise TypeError(f'{name} got incompatible shapes for broadcasting: '
   1618                       f'{", ".join(map(str, map(tuple, shapes)))}.')
   1620 return tuple(result_shape)

TypeError: mul got incompatible shapes for broadcasting: (4,), (3,).
```


Input type for deterministic
It seems like the expected input type for `deterministic()` is `numpy.ndarray`. This causes my IDE (PyCharm) to complain when I pass in a JAX `Array` type. Since deterministic variables are often the result of applying JAX functions, I would think that `Array` types would be expected.


Discussion:

Resolving patch:
diff --git a/numpyro/distributions/transforms.py b/numpyro/distributions/transforms.py
--- a/numpyro/distributions/transforms.py
+++ b/numpyro/distributions/transforms.py
@@ -871,6 +871,12 @@ def log_abs_det_jacobian(self, x, y, intermediates=None):
         J_logdet = (softplus(y) + softplus(-y)).sum(-1)
         return J_logdet
 
+    def forward_shape(self, shape):
+        return shape[:-1] + (shape[-1] - 1,)
+
+    def inverse_shape(self, shape):
+        return shape[:-1] + (shape[-1] + 1,)
+
 
 def _softplus_inv(y):
     return jnp.log(-jnp.expm1(-y)) + y
diff --git a/numpyro/primitives.py b/numpyro/primitives.py
--- a/numpyro/primitives.py
+++ b/numpyro/primitives.py
@@ -135,14 +135,14 @@ def sample(
 
     :param str name: name of the sample site.
     :param fn: a stochastic function that returns a sample.
-    :param numpy.ndarray obs: observed value
+    :param jnp.ndarray obs: observed value
     :param jax.random.PRNGKey rng_key: an optional random key for `fn`.
     :param sample_shape: Shape of samples to be drawn.
     :param dict infer: an optional dictionary containing additional information
         for inference algorithms. For example, if `fn` is a discrete distribution,
         setting `infer={'enumerate': 'parallel'}` to tell MCMC marginalize
         this discrete latent site.
-    :param numpy.ndarray obs_mask: Optional boolean array mask of shape
+    :param jnp.ndarray obs_mask: Optional boolean array mask of shape
         broadcastable with ``fn.batch_shape``. If provided, events with
         mask=True will be conditioned on ``obs`` and remaining events will be
         imputed by sampling. This introduces a latent sample site named ``name
@@ -235,7 +235,7 @@ def param(name, init_value=None, **kwargs):
         Note that the onus of using this to initialize the optimizer is
         on the user inference algorithm, since there is no global parameter
         store in NumPyro.
-    :type init_value: numpy.ndarray or callable
+    :type init_value: jnp.ndarray or callable
     :param constraint: NumPyro constraint, defaults to ``constraints.real``.
     :type constraint: numpyro.distributions.constraints.Constraint
     :param int event_dim: (optional) number of rightmost dimensions unrelated
@@ -289,7 +289,7 @@ def deterministic(name, value):
     values in the model execution trace.
 
     :param str name: name of the deterministic site.
-    :param numpy.ndarray value: deterministic value to record in the trace.
+    :param jnp.ndarray value: deterministic value to record in the trace.
     """
     if not _PYRO_STACK:
         return value
@@ -376,7 +376,7 @@ def model():
             # ...
 
     :returns: The mask.
-    :rtype: None, bool, or numpy.ndarray
+    :rtype: None, bool, or jnp.ndarray
     """
     return _inspect()["mask"]
 
@@ -607,7 +607,7 @@ def factor(name, log_factor):
     probabilistic model.
 
     :param str name: Name of the trivial sample.
-    :param numpy.ndarray log_factor: A possibly batched log probability factor.
+    :param jnp.ndarray log_factor: A possibly batched log probability factor.
     """
     unit_dist = numpyro.distributions.distribution.Unit(log_factor)
     unit_value = unit_dist.sample(None)
@@ -657,11 +657,11 @@ def model(data):
                 data = numpyro.subsample(data, event_dim=0)
                 # ...
 
-    :param numpy.ndarray data: A tensor of batched data.
+    :param jnp.ndarray data: A tensor of batched data.
     :param int event_dim: The event dimension of the data tensor. Dimensions to
         the left are considered batch dimensions.
     :returns: A subsampled version of ``data``
-    :rtype: ~numpy.ndarray
+    :rtype: ~jnp.ndarray
     """
     if not _PYRO_STACK:
         return data


Test patch:
diff --git a/test/test_distributions.py b/test/test_distributions.py
--- a/test/test_distributions.py
+++ b/test/test_distributions.py
@@ -2300,6 +2300,16 @@ def test_composed_transform_1(batch_shape):
     assert_allclose(log_det, expected_log_det)
 
 
+@pytest.mark.parametrize("batch_shape", [(), (5,)])
+def test_simplex_to_order_transform(batch_shape):
+    simplex = jnp.arange(5.0) / jnp.arange(5.0).sum()
+    simplex = jnp.broadcast_to(simplex, batch_shape + simplex.shape)
+    transform = SimplexToOrderedTransform()
+    out = transform(simplex)
+    assert out.shape == transform.forward_shape(simplex.shape)
+    assert simplex.shape == transform.inverse_shape(out.shape)
+
+
 @pytest.mark.parametrize("batch_shape", [(), (5,)])
 @pytest.mark.parametrize("prepend_event_shape", [(), (4,)])
 @pytest.mark.parametrize("sample_shape", [(), (7,)])