internlm/SWE-Fixer-Train-110K huggingface dataset:
Repository: pyro-ppl/numpyro

Issue:
Raise an exception (or warning) when calling the inverse of a non-invertible transform.
The `Transform` class in `numpyro.distributions.transforms` provides an `inv` property which is meant to return the inverse transform. However, this is available even for functions that are not invertible. In those cases, we silently return the identity function which seems weird.

### Example:
```python
import numpyro.distributions as dist

f1 = dist.transforms.AbsTransform()
print(f1.inv(-1))       # Not defined, but returns -1
print(f1.inv(f1(-1))) # Returns +1
```

The issue is also present in transformations that are obtained via `ComposeTransform`.
When the transformation is invertible, it is very easy to obtain the CDF of the transformed distribution (which is something I often need). So I think providing an 'automatic' inverse function is a good idea, but maybe it'd be good to raise an error (or at least a warning) when trying to access the inverse of a non-invertible transform. 


Discussion:
Yeah, I think it is better to raise a warning. (I'm not sure if we want to take inverse of abs transform some where)
@omarfsosa Do you want to add that warning?

Resolving patch:
diff --git a/numpyro/distributions/transforms.py b/numpyro/distributions/transforms.py
--- a/numpyro/distributions/transforms.py
+++ b/numpyro/distributions/transforms.py
@@ -160,6 +160,11 @@ def __call__(self, x):
         return jnp.abs(x)
 
     def _inverse(self, y):
+        warnings.warn(
+            "AbsTransform is not a bijective transform."
+            " The inverse of `y` will be `y`.",
+            stacklevel=find_stack_level(),
+        )
         return y
 
 
diff --git a/numpyro/infer/reparam.py b/numpyro/infer/reparam.py
--- a/numpyro/infer/reparam.py
+++ b/numpyro/infer/reparam.py
@@ -187,7 +187,7 @@ class NeuTraReparam(Reparam):
     """
     Neural Transport reparameterizer [1] of multiple latent variables.
 
-    This uses a trained :class:`~pyro.contrib.autoguide.AutoContinuous`
+    This uses a trained :class:`~numpyro.infer.autoguide.AutoContinuous`
     guide to alter the geometry of a model, typically for use e.g. in MCMC.
     Example usage::
 
diff --git a/numpyro/util.py b/numpyro/util.py
--- a/numpyro/util.py
+++ b/numpyro/util.py
@@ -444,9 +444,9 @@ def format_shapes(
         def model(*args, **kwargs):
             ...
 
-        with numpyro.handlers.seed(rng_key=1):
+        with numpyro.handlers.seed(rng_seed=1):
             trace = numpyro.handlers.trace(model).get_trace(*args, **kwargs)
-        numpyro.util.format_shapes(trace)
+        print(numpyro.util.format_shapes(trace))
     """
     if not trace.keys():
         return title


Test patch:
Repository: pyro-ppl/numpyro

Issue:
`numpyro.render_model` behaviour changed in 0.9.0
I noticed that the behaviour of `numpyro.render_model` appears to have changed since the latest release, specifically arrows signifying conditional dependence are not being rendered properly.

I haven't had time to dig into why this is yet, but here is a simple recreation

```python
import numpyro
import numpyro.distributions as dist


def model():
    theta = numpyro.sample("theta", dist.Uniform())
    
    with numpyro.plate("N", size=10):
        numpyro.sample("y", dist.Bernoulli(theta))
        
        
if __name__ == "__main__":
    g = numpyro.render_model(model)
    g.render(f"model-{numpyro.__version__}", format="png")
```

Running this with 0.8.0 and 0.9.0 I get (respectively)

![model-0 8 0](https://user-images.githubusercontent.com/15220906/152194532-07f70d39-8197-41c5-9e14-beefc33b807e.png)

![model-0 9 0](https://user-images.githubusercontent.com/15220906/152194539-98a41b9f-a673-4937-b864-1fb9e6c3d420.png)




Discussion:
Seems like it may be related to #1248?
Something has changed in `get_model_relations` in `numpyro.contrib.render`. I get different outputs there with 0.8.0 and 0.9.0

```diff
{
- 'sample_sample': {'theta': [], 'y': ['theta']}, 
+ 'sample_sample': {'theta': [], 'y': []}, 
  'sample_dist': {'theta': 'Uniform', 'y': 'BernoulliProbs'}, 
  'plate_sample': {'N': ['y']}, 
  'observed': [],
}
```
Could you try to use `jax 0.2.27`? I suspected the new release caused the problem (due to a recent support for dynamic shapes: https://github.com/google/jax/pull/9316). We need to update Provenance logic.
Yep! Downgrading JAX resolves the issue.

Resolving patch:
diff --git a/numpyro/ops/provenance.py b/numpyro/ops/provenance.py
--- a/numpyro/ops/provenance.py
+++ b/numpyro/ops/provenance.py
@@ -26,6 +26,9 @@ def process_primitive(self, primitive, tracers, params):
             out_tracers = out_tracers if primitive.multiple_results else [out_tracers]
             for t in out_tracers:
                 t.aval.named_shape["_provenance"] = out_provenance
+                # Also update provenance of the cached tracer -> aval dict.
+                aval_cache = self.frame.tracer_to_var[id(t)].aval
+                aval_cache.named_shape["_provenance"] = out_provenance
             out_tracers = out_tracers if primitive.multiple_results else out_tracers[0]
         return out_tracers
 
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -9,7 +9,7 @@
 from setuptools import find_packages, setup
 
 PROJECT_PATH = os.path.dirname(os.path.abspath(__file__))
-_jax_version_constraints = ">=0.2.13,<0.2.28"
+_jax_version_constraints = ">=0.2.13"
 _jaxlib_version_constraints = ">=0.1.65"
 
 # Find version


Test patch:
Repository: pyro-ppl/numpyro

Issue:
compat module: compat.infer.MCMC broken
Calling
```python
compat.infer.MCMC(kernel, num_warmup=100, num_samples=1000)
```
throws an error since no kwargs are handed to `mcmc.MCMC`. Also, shouldnt the correct `compat` call be
```python
compat.infer.MCMC(kernel, warmup_steps=100, num_samples=1000)
```
Thanks :)


Discussion:
Thanks, @alexander-maier! I think we'll need to use [keyword arguments here](https://github.com/pyro-ppl/numpyro/blob/9719466a584af80683a7aa04401d4e065cc569ae/numpyro/compat/infer.py#L99-L100). It is surprised to me that we don't have a test for this.

Resolving patch:
diff --git a/numpyro/compat/infer.py b/numpyro/compat/infer.py
--- a/numpyro/compat/infer.py
+++ b/numpyro/compat/infer.py
@@ -82,7 +82,7 @@ def __init__(
         self,
         kernel,
         num_samples,
-        num_warmup=None,
+        warmup_steps=None,
         initial_params=None,
         num_chains=1,
         hook_fn=None,
@@ -91,13 +91,13 @@ def __init__(
         disable_validation=True,
         transforms=None,
     ):
-        if num_warmup is None:
-            num_warmup = num_samples
+        if warmup_steps is None:
+            warmup_steps = num_samples
         self._initial_params = initial_params
         self._mcmc = mcmc.MCMC(
             kernel,
-            num_warmup,
-            num_samples,
+            num_warmup=warmup_steps,
+            num_samples=num_samples,
             num_chains=num_chains,
             progress_bar=(not disable_progbar),
         )


Test patch:
Repository: pyro-ppl/numpyro

Issue:
0.9.0 error when sampling posterior predictive via guide
Error:
`TypeError: mul got incompatible shapes for broadcasting: (1000,), (3,).`

Works just fine in `0.8.0`

Code to reproduce:
```python
import warnings

warnings.filterwarnings("ignore", category=DeprecationWarning)

import pandas as pd

import arviz as az
import matplotlib.pyplot as plt

import jax.numpy as jnp
from jax import random

import numpyro
import numpyro.distributions as dist
import numpyro.optim as optim
from numpyro.infer import SVI, Trace_ELBO, Predictive
from numpyro.infer.autoguide import AutoLaplaceApproximation, AutoNormal


# INFO: load data for problems from the official repo
data_uri = "https://raw.githubusercontent.com/rmcelreath/rethinking/master/data/Howell1.csv"
df_dev = pd.read_csv(data_uri, sep=";")
df_dev = df_dev[df_dev["age"] >= 18]

X_train = df_dev.loc[:, ["height"]]
y_train = df_dev.loc[:, ["weight"]]

df_test = pd.DataFrame(data={"height": [140, 160, 175], "weight": [None, None, None]})
X_test = df_test.loc[:, ["height"]]
y_test = df_test.loc[:, ["weight"]]


def model(X: pd.DataFrame, y: pd.DataFrame, observed=True):
    alpha = numpyro.sample("alpha", dist.Normal(60, 10))
    beta = numpyro.sample("beta", dist.LogNormal(0, 1))
    sigma = numpyro.sample("sigma", dist.Uniform(0, 10))
    # INFO: subtract X_train explicitly to secure non-leaky inference
    mu = numpyro.deterministic("mu", alpha + beta * (X["height"] - X_train["height"].mean()).values)
    numpyro.sample("W", dist.Normal(mu, sigma), obs=y["weight"].values if observed else None)


# quadratic approximation part
guide = AutoNormal(model)
svi = SVI(model, guide, optim.Adam(1), Trace_ELBO(), X=X_train, y=y_train)
svi_result = svi.run(random.PRNGKey(0), 1000)

params = svi_result.params

# ---

# get posterior predictive samples, approach 1: from samples
num_samples = 1000
samples_posterior = guide.sample_posterior(random.PRNGKey(1), params=params, sample_shape=(num_samples,))
dist_posterior_predictive = Predictive(model=model, guide=guide, params=params, posterior_samples=samples_posterior)
samples_posterior_predictive = dist_posterior_predictive(random.PRNGKey(1), X=X_test, y=None, observed=False)
samples_posterior_predictive

```


Discussion:
Thanks, @ColdTeapot273K! We don't support providing both posterior_samples and guide. Will make a PR to raise the error soon.

Resolving patch:
diff --git a/numpyro/infer/util.py b/numpyro/infer/util.py
--- a/numpyro/infer/util.py
+++ b/numpyro/infer/util.py
@@ -862,6 +862,10 @@ def __init__(
             raise ValueError(
                 "Either posterior_samples or num_samples must be specified."
             )
+        if posterior_samples is not None and guide is not None:
+            raise ValueError(
+                "Only one of guide or posterior_samples can be provided, not both."
+            )
 
         batch_ndims = (
             batch_ndims if batch_ndims is not None else 1 if guide is None else 0


Test patch:
Repository: pyro-ppl/numpyro

Issue:
Add option to exclude `plate` from `scope` handler
Currently, `scope` handler changes the name of all messages. I found that it will cause a bit confusing so it would be nice to have an option to exclude some message types. In the trace, we only raise issues when `sample` and `deterministic` messages are duplicated, so it makes sense to only rename them by default.


Discussion:

Resolving patch:
diff --git a/numpyro/contrib/control_flow/scan.py b/numpyro/contrib/control_flow/scan.py
--- a/numpyro/contrib/control_flow/scan.py
+++ b/numpyro/contrib/control_flow/scan.py
@@ -153,7 +153,9 @@ def body_fn(wrapped_carry, x, prefix=None):
 
         if init:
             # handler the name to match the pattern of sakkar_bilmes product
-            with handlers.scope(prefix="_PREV_" * (unroll_steps - i), divider=""):
+            with handlers.scope(
+                prefix="_PREV_" * (unroll_steps - i), divider="", hide_types=["plate"]
+            ):
                 new_carry, y = config_enumerate(seeded_fn)(carry, x)
                 trace = {}
         else:
diff --git a/numpyro/handlers.py b/numpyro/handlers.py
--- a/numpyro/handlers.py
+++ b/numpyro/handlers.py
@@ -625,18 +625,20 @@ class scope(Messenger):
     :param fn: Python callable with NumPyro primitives.
     :param str prefix: a string to prepend to sample names
     :param str divider: a string to join the prefix and sample name; default to `'/'`
+    :param list hide_types: an optional list of side types to skip renaming.
     """
 
-    def __init__(self, fn=None, prefix="", divider="/"):
+    def __init__(self, fn=None, prefix="", divider="/", *, hide_types=None):
         self.prefix = prefix
         self.divider = divider
+        self.hide_types = [] if hide_types is None else hide_types
         super().__init__(fn)
 
     def process_message(self, msg):
-        if msg.get("name"):
+        if msg.get("name") and msg["type"] not in self.hide_types:
             msg["name"] = f"{self.prefix}{self.divider}{msg['name']}"
 
-        if msg.get("cond_indep_stack"):
+        if msg.get("cond_indep_stack") and "plate" not in self.hide_types:
             msg["cond_indep_stack"] = [
                 CondIndepStackFrame(
                     f"{self.prefix}{self.divider}{i.name}", i.dim, i.size


Test patch:
diff --git a/test/test_handlers.py b/test/test_handlers.py
--- a/test/test_handlers.py
+++ b/test/test_handlers.py
@@ -577,18 +577,18 @@ def test_block():
 
 def test_scope():
     def fn():
-        return numpyro.sample("x", dist.Normal())
+        with numpyro.plate("N", 10):
+            return numpyro.sample("x", dist.Normal())
 
     with handlers.trace() as trace:
         with handlers.seed(rng_seed=1):
             with handlers.scope(prefix="a"):
                 fn()
             with handlers.scope(prefix="b"):
-                with handlers.scope(prefix="a"):
+                with handlers.scope(prefix="a", hide_types=["plate"]):
                     fn()
 
-    assert "a/x" in trace
-    assert "b/a/x" in trace
+    assert set(trace) == {"a/x", "b/a/x", "a/N", "b/N"}
 
 
 def test_scope_frames():
Repository: pyro-ppl/numpyro

Issue:
signature error when running SA MCMC sampler in parallel
the following code fails:

```
from jax.random import PRNGKey

import numpyro
import numpyro.distributions as dist
import numpyro.infer as infer

numpyro.set_platform('cpu')
numpyro.set_host_device_count(4)


def model():
    y = numpyro.sample("y", dist.Normal(0.0, 1.0))


kernel = infer.SA(model) 
#kernel = infer.HMC(model)  # works

mcmc = infer.MCMC(
    kernel,
    num_warmup=0,
    num_samples=100,
    num_chains=4,
    #num_chains=1,  # works
    #chain_method='sequential',  # works
)


mcmc.run(PRNGKey(0))
```
with the error (and a long stack trace):
```
TypeError: body_fun output and input must have same type structure, got PyTreeDef((*, *, ({'y': *}, *, None), *)) and PyTreeDef((*, *, ({'y': *}, *, {'y': *}), *)).
```
any of the annotated changes make the code work again. what gives?


Discussion:

Resolving patch:
diff --git a/numpyro/infer/util.py b/numpyro/infer/util.py
--- a/numpyro/infer/util.py
+++ b/numpyro/infer/util.py
@@ -374,7 +374,8 @@ def body_fn(state):
         return i + 1, key, (params, pe, z_grad), is_valid
 
     def _find_valid_params(rng_key, exit_early=False):
-        init_state = (0, rng_key, (prototype_params, 0.0, prototype_params), False)
+        prototype_grads = prototype_params if validate_grad else None
+        init_state = (0, rng_key, (prototype_params, 0.0, prototype_grads), False)
         if exit_early and not_jax_tracer(rng_key):
             # Early return if valid params found. This is only helpful for single chain,
             # where we can avoid compiling body_fn in while_loop.


Test patch:
diff --git a/test/infer/test_mcmc.py b/test/infer/test_mcmc.py
--- a/test/infer/test_mcmc.py
+++ b/test/infer/test_mcmc.py
@@ -869,13 +869,15 @@ def model():
     mcmc.run(random.PRNGKey(0))
 
 
-def test_SA_gradient_free():
+@pytest.mark.parametrize("num_chains", [1, 2])
+@pytest.mark.filterwarnings("ignore:There are not enough devices:UserWarning")
+def test_SA_chain_gradient_free(num_chains):
     def model():
         x = numpyro.sample("x", dist.Normal(0, 1))
         y = lax.while_loop(lambda x: x < 10, lambda x: x + 1, x)
         numpyro.sample("obs", dist.Normal(y, 1), obs=1.0)
 
-    mcmc = MCMC(SA(model), num_warmup=10, num_samples=10)
+    mcmc = MCMC(SA(model), num_warmup=10, num_samples=10, num_chains=num_chains)
     mcmc.run(random.PRNGKey(0))
Repository: pyro-ppl/numpyro

Issue:
issues when numpyro.enable_validation(is_validate=True)
Hi
I identified a strange behavior: 
I estimate a VectorAutoRegression model and when  I use numpyro.enable_validation(is_validate=True), (which I had in the code
for some versions of the model with a MultivariateTruncated) , the resuts of any simple version of the model (with MultivariateNormal residuals say) were off. 
Taking out numpyro.enable_validation(is_validate=True) the results look fine.
Should I be worried?

best




Discussion:
I think there might be some problems. When enabling the validation, we mask out-of-support samples by zero density. Not sure if it is related. Could you make reproducible code?


Ok here is an example, if you comment line 11 it returns sensible results. Otherwise it fails. The data are downloaded via pandasdmx. If you prefer I can send you a csv.


@hyperfra downloading zip files from unknown third parties isn't exactly a great idea. can you please instead provide e.g. a gist and create fake data within the script?
sorry, you are quite right. now the link works
https://gist.github.com/hyperfra/9e8df68414b0172c852a16db82659bc6


Thanks @hyperfra! It seems that there is something wrong with the validation logic of LKJCholesky distribution.

Resolving patch:
diff --git a/numpyro/distributions/constraints.py b/numpyro/distributions/constraints.py
--- a/numpyro/distributions/constraints.py
+++ b/numpyro/distributions/constraints.py
@@ -115,7 +115,8 @@ def __call__(self, x):
         )
         positive_diagonal = jnp.all(jnp.diagonal(x, axis1=-2, axis2=-1) > 0, axis=-1)
         x_norm = jnp.linalg.norm(x, axis=-1)
-        unit_norm_row = jnp.all((x_norm <= 1) & (x_norm > 1 - 1e-6), axis=-1)
+        tol = jnp.finfo(x.dtype).eps * x.shape[-1] * 10
+        unit_norm_row = jnp.all(jnp.abs(x_norm - 1) <= tol, axis=-1)
         return lower_triangular & positive_diagonal & unit_norm_row
 
     def feasible_like(self, prototype):


Test patch:
Repository: pyro-ppl/numpyro

Issue:
Mixture: Intermediates cannot be coerced to bool
```python
from jax import numpy as jnp
from jax.random import PRNGKey
from numpyro import distributions as dist
from numpyro import sample
from numpyro import handlers as hdl
from numpyro.infer.util import log_density

key = PRNGKey(0)

def model(toggle):
    d1 = dist.HalfNormal()
    d2 = dist.LogNormal()
    mixing = dist.Categorical(jnp.full((2,), 0.5))
    mixture = dist.Mixture(mixing, [d1, d2])
    if toggle:
        sample('s', mixture)
    else:
        sample('s', mixture, sample_shape=(2,))

with hdl.seed(rng_seed=key):
    l1 = log_density(model, [], {'toggle': True}, {})
    l2 = log_density(model, [], {'toggle': False}, {})
```

In the above `l1` evaluates fine, as the intermediates returned in the model trace is a `DeviceArray` containing a single number, but when sample/batch shape is anything but `(1,)`, `intermediates` is an array with size > 1 the following boolean coercion causes and error: 

![image](https://user-images.githubusercontent.com/392251/197748939-3dfa2096-4d2a-4451-9371-d8187bfcefee.png)

Comparing `_MixtureBase` to other distributions with `sample_with_intermediates`, it seems all other cases return a list of intermediates, whereas `_MixtureBase` simple returns a jax array directly. It appears the most logical fix would simply be `_MixtureBase.sample_with_intermediates` returning the intermediates array wrapped a list in order to comply with the current API. If this sounds like a reasonable fix, I would be happy to send a PR.


Discussion:

Resolving patch:
diff --git a/numpyro/distributions/mixtures.py b/numpyro/distributions/mixtures.py
--- a/numpyro/distributions/mixtures.py
+++ b/numpyro/distributions/mixtures.py
@@ -139,7 +139,7 @@ def sample_with_intermediates(self, key, sample_shape=()):
         )
 
         # Final sample shape (*sample_shape, *batch_shape, *event_shape)
-        return jnp.squeeze(samples_selected, axis=self.mixture_dim), indices
+        return jnp.squeeze(samples_selected, axis=self.mixture_dim), [indices]
 
     def sample(self, key, sample_shape=()):
         return self.sample_with_intermediates(key=key, sample_shape=sample_shape)[0]
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -51,6 +51,7 @@
         "test": [
             "black[jupyter]>=21.8b0",
             "flake8",
+            "importlib-metadata<5.0",
             "isort>=5.0",
             "pytest>=4.1",
             "pyro-api>=0.1.1",


Test patch:
diff --git a/test/test_distributions_mixture.py b/test/test_distributions_mixture.py
--- a/test/test_distributions_mixture.py
+++ b/test/test_distributions_mixture.py
@@ -115,7 +115,7 @@ def _test_mixture(mixing_distribution, component_distribution):
     expected_shape = samples.shape[:nb_value_dims]
     assert lp.shape == expected_shape
     # Samples with indices
-    samples_, indices = mixture.sample_with_intermediates(
+    samples_, [indices] = mixture.sample_with_intermediates(
         rng_key, sample_shape=sample_shape
     )
     assert samples_.shape == samples.shape
Repository: pyro-ppl/numpyro

Issue:
`MixtureSameFamily` outputting incorrect log probs
```python
def gmm_log_prob(N=5000):
    mix = Categorical(logits=jnp.ones((N, 2)),)
    comp = Normal(jnp.zeros((N, 2)), jnp.ones((N,2)),)
    gmm = MixtureSameFamily(mix, comp,)
    data = jnp.zeros((N,))
    return gmm.log_prob(data),  Normal(0, 1).log_prob(data)
```

```
gmm_log_prob()
(DeviceArray([0.77420866, 0.77420866, 0.77420866, ..., 0.77420866,
              0.77420866, 0.77420866], dtype=float32),
 DeviceArray([-0.9189385, -0.9189385, -0.9189385, ..., -0.9189385,
              -0.9189385, -0.9189385], dtype=float32))
```

Distrax and torch provide the correct output, which should be equal here


Discussion:
Thanks for catching this! It looks like we didn't normalize the logits [here](https://github.com/pyro-ppl/numpyro/blob/f81ce325de402d9a9891497136a36e3b44f982cc/numpyro/distributions/mixtures.py#L275) and [here](https://github.com/pyro-ppl/numpyro/blob/f81ce325de402d9a9891497136a36e3b44f982cc/numpyro/distributions/mixtures.py#L436). Do you want to submit a fix? something like
```
     return jax.nn.log_softmax(self.mixing_distribution.logits) + component_log_probs
```

Resolving patch:
diff --git a/numpyro/distributions/mixtures.py b/numpyro/distributions/mixtures.py
--- a/numpyro/distributions/mixtures.py
+++ b/numpyro/distributions/mixtures.py
@@ -272,7 +272,7 @@ def component_sample(self, key, sample_shape=()):
     def component_log_probs(self, value):
         value = jnp.expand_dims(value, self.mixture_dim)
         component_log_probs = self.component_distribution.log_prob(value)
-        return self.mixing_distribution.logits + component_log_probs
+        return jax.nn.log_softmax(self.mixing_distribution.logits) + component_log_probs
 
 
 class MixtureGeneral(_MixtureBase):
@@ -433,7 +433,7 @@ def component_log_probs(self, value):
         component_log_probs = jnp.stack(
             [d.log_prob(value) for d in self.component_distributions], axis=-1
         )
-        return self.mixing_distribution.logits + component_log_probs
+        return jax.nn.log_softmax(self.mixing_distribution.logits) + component_log_probs
 
 
 def _check_mixing_distribution(mixing_distribution):


Test patch:
diff --git a/test/test_distributions.py b/test/test_distributions.py
--- a/test/test_distributions.py
+++ b/test/test_distributions.py
@@ -1182,6 +1182,15 @@ def test_log_prob(jax_dist, sp_dist, params, prepend_shape, jit):
     assert_allclose(jit_fn(jax_dist.log_prob)(samples), expected, atol=1e-5)
 
 
+def test_mixture_log_prob():
+    gmm = dist.MixtureSameFamily(
+        dist.Categorical(logits=np.zeros(2)), dist.Normal(0, 1).expand([2])
+    )
+    actual = gmm.log_prob(0.0)
+    expected = dist.Normal(0, 1).log_prob(0.0)
+    assert_allclose(actual, expected)
+
+
 @pytest.mark.parametrize(
     "jax_dist, sp_dist, params",
     # TODO: add more complete pattern for Discrete.cdf