Repository: pyro-ppl/numpyro

Issue:
How to use vmap in the context of a model?
I was excited to see that numpyro now supports vmap. But I'm a little confused about what exactly can be done with it. When I do this:

```
import numpyro
import jax
from jax import numpy as jnp
dist = numpyro.distributions

loc = jnp.array([0,1])
scale = jnp.array([0,1])
d = jax.vmap(dist.Normal, in_axes=(0,0))(loc,scale)
key = jax.random.PRNGKey(0)
d.sample(key)
```

Then it works fine—I get a sample from a standard normal in two dimensions. However, when I try to do this:

```
import numpyro
import jax
from jax import numpy as jnp
dist = numpyro.distributions

loc = jnp.array([0,1])
scale = jnp.array([0,1])
def model():
    numpyro.sample("x",jax.vmap(dist.Normal, in_axes=(0,0))(loc,scale))

kernel = numpyro.infer.NUTS(model)
mcmc = numpyro.infer.MCMC(kernel, num_warmup=250, num_samples=750)
key = jax.random.PRNGKey(0)
mcmc.run(key)
```

Then I get `TypeError: '>' not supported between instances of 'object' and 'float'` which seem to be an error in `_GreaterThan.__call__` in `numpyro/distributions/constraints.py`.

Am I using `vmap` incorrectly here? Is there any other way that I can `vmap` distributions and then build numpyro models that I can do MCMC on, etc.?


Discussion:
I think this has something to do with validation.
```
with dist.distribution.validation_enabled():
    jax.vmap(dist.Normal)(loc, scale)
```
raises the above error.

@pierreglaser do you want to take a look?
The problem comes from `tree_unflatten` calling `Distribution.__init__` when`Distribution._validate_args` is set to `True`, something which is not tested right now, hence the bug. @fehiepsi, WDYT about turning of validation in `tree_unflatten`? This could be done by skipping the `Distribution.__init__` call, or creating a `validation_disabled` context manager to override enclosing validation instructions?
Hi @pierreglaser, setting `validate_args=False` in that __init__ call makes sense to me. We might want to override `_validate_args` after that call to respect the global value.
Thank you both for the fast response! Happy to wait for an update, but in the meantime is there by any chance some kind of workaround? E.g. might there by some kind of magic function like `numpyro_config` that would allow me to do this?

```
def model():
    with numpyro_config(validate_args=False):
        numpyro.sample("x",jax.vmap(dist.Normal, in_axes=(0,0))(loc,scale))
```

Or something like that?
@justindomke a couple clarifications regarding `numpyro` and `vmap`.

First, let me clarify your claim `numpyro` now supports `vmap`: the only thing that the recent feature added by #1529  is that it made `Distribution` compliant [`pytrees`](https://jax.readthedocs.io/en/latest/pytrees.html), with all the benefits that comes with it. In particular, this means that `Distribution` objects can be used as inputs and outputs of functions transformed by jax program transformations, like `jit`, `jvp`, and in particular, `vmap`. However, this feature does not allow you to plug in `vmap`-ed `Distributions` object in any place where a standard distribution is expected, something that is out of the scope of the `vmap` model. 

Keeping that in mind, your first example: 

```
import numpyro
import jax
from jax import numpy as jnp
dist = numpyro.distributions

loc = jnp.array([0,1])
scale = jnp.array([0,1])
d = jax.vmap(dist.Normal, in_axes=(0,0))(loc,scale)
key = jax.random.PRNGKey(0)
d.sample(key)
```

assumes that d, a **`vmap`-ed** `Distribution` object, implements a `sample` method with the same semantics than the `sample` method non `vmap`-ed `Distribution`. But as explained above, this call is not guaranteed  to work as you expect it to work: under the `vma`p model,  A vmap-ed object `v_obj` of a class `Obj` defining a method `Obj.f(self):...` does not support calls of the form `v_obj.f()` Instead, it supports calls of the form `vmap(Obj.f)(v_obj)`. 
Thus, although your first code snippets gives you a result, it is purely coincidental, and not guaranteed in other cases of more elaborate mapping axes, or other `Distributions` classes.

Instead, a usage of `vmap` and `numpyro` with guaranteed support that achieves a similar functionality as your first example would be:
```python
import numpyro
import jax
from jax import numpy as jnp
from numpyro.distributions.batch_util import vmap_over
dist = numpyro.distributions


loc = jnp.array([0,1])
scale = jnp.array([0,1])
v_dist = jax.vmap(dist.Normal, in_axes=(0, 0))(loc, scale)

key = jax.random.PRNGKey(0)
samples = jax.vmap(dist.Normal.sample, in_axes=(vmap_over(v_dist, loc=0, scale=0), 0))(v_dist, jax.random.split(key))
```

Obviously, this use of `numpyro` and `vmap` is quite simple. More complicated settings leveraging that feature include batched MCMC sampling of multiple intractable `Distribution` instances, as described at length in [this thread](https://github.com/pyro-ppl/numpyro/pull/1529#issuecomment-1411403617). 

Similarly, your second example

```python
import numpyro
import jax
from jax import numpy as jnp
dist = numpyro.distributions

loc = jnp.array([0,1])
scale = jnp.array([0,1])
def model():
    numpyro.sample("x",jax.vmap(dist.Normal, in_axes=(0,0))(loc,scale))

kernel = numpyro.infer.NUTS(model)
mcmc = numpyro.infer.MCMC(kernel, num_warmup=250, num_samples=750)
key = jax.random.PRNGKey(0)
mcmc.run(key)
```

makes a similar substitution-ability assumption between `vmap`-ed and non-`vmap`-ed `Distribution` objects, and is thus likely to error out. A snippet that achieves a similar functionality and that leverage the new `Distribution`-as-pytree feature recently introduced is:
```python
import numpyro
import jax
from jax import numpy as jnp
dist = numpyro.distributions


def run_mcmc(d, key):
    def model():
        numpyro.sample("x", d)
    kernel = numpyro.infer.NUTS(model)
    mcmc = numpyro.infer.MCMC(kernel, num_warmup=250, num_samples=750)
    mcmc.run(key)

v_dist = jax.vmap(dist.Normal, in_axes=(0, 0))(jnp.zeros((2,)), jnp.zeros((2,)))
key = jax.random.PRNGKey(0)
jax.vmap(run_mcmc, in_axes=(0, 0))(v_dist, jax.random.split(key))
```

notice how the vmapped distribution is passed as an argument of a transformed version of a function expecting a standard distribution, it is not passed directly to `run_mcmc`. Without this feature, one could not have had used such a function signature, since `Distributions` objects were not compliant pytrees, and thus could not be passed to  functions to be transformed by jax. Instead, one would have had unpack the `Distribution` individual arguments:

```python
import numpyro
import jax
from jax import numpy as jnp
dist = numpyro.distributions


loc = jnp.array([0,1])
scale = jnp.array([0,1])
def run_mcmc(key, loc, scale):
    def model():
        numpyro.sample("x",dist.Normal(loc, scale))
    kernel = numpyro.infer.NUTS(model)
    mcmc = numpyro.infer.MCMC(kernel, num_warmup=250, num_samples=750)
    mcmc.run(key)

key = jax.random.PRNGKey(0)
jax.vmap(run_mcmc, in_axes=(0, 0, 0))(jax.random.split(key), loc, scale)
```

which is much clumsier, which is why this feature is beneficial.

Finally, it seems that the last 2 code snippts (using `MCMC` with or without #1529) actually error out 😄. We should fix that @fehiepsi. However, the reason seems orthogonal to #1529. 

@pierreglaser thanks very much for the clarification! Regarding the second example, I see how your fix works, but I was hoping to be able to address a more general case where there are upstream random variables. A more realistic example would be something like this:

import numpyro
import jax
from jax import numpy as jnp
dist = numpyro.distributions

```python
scale = jnp.array([0,1])
def model():
    z = numpyro.sample("z",dist.Normal(0,1))
    x = numpyro.sample("x",jax.vmap(dist.Normal, in_axes=(None,0))(z,scale))
```
```python
kernel = numpyro.infer.NUTS(model)
mcmc = numpyro.infer.MCMC(kernel, num_warmup=250, num_samples=750)
key = jax.random.PRNGKey(0)
mcmc.run(key)
```

Here the two different dimensions of `x` are tied together through the latent variable `z`. So doing a vmap over the full run of MCMC wouldn't be helpful here.

I know that `plate` can be used for things like this, but vmap is more powerful so I was hoping to be able to use it instead. But if this use-case is totally outside of what's intended to work, that's understandable!
```python
scale = jnp.array([0,1])
def model():
    z = numpyro.sample("z",dist.Normal(0,1))
    x = numpyro.sample("x",jax.vmap(dist.Normal, in_axes=(None,0))(z,scale))
```

`plate`may work, or even just using `dist.Normal(z, scale)` and let  `numpyro`'s shape broadcasting mechanism do its thing right? 

> But if this use-case is totally outside of what's intended to work, that's understandable!

So far it is. In general, I don' think arbitrary substitionability between `Distributions` and `vmap`-ed `Distributions` is possible, or at least without a lot of auto-magic heavy lifting that is outside the very scopes where `jax` advertises the uses of `vmap`. 
Of course, this particular example is easy to do with plate or auto broadcasting (or even a for loop) but there are more complex cases that are quite tricky and `jax.vmap` would be helpful.

If you're interested, the *real* reason I was asking is because I'm developing another PPL (https://github.com/justindomke/pangolin) that supports multiple inference backends. I already have pretty full support for numpyro as a backend with arbitrary uses of vmap in defining models, but I do it by sort of building my own potential function that vmaps over calls to things like `Normal(loc,scale).log_prob(z)` and whatnot and then calling numpyro's inference routines on that potential function, rather than using numpyro's abstraction of "models". This works fine, but I do lose some of numpyro's goodies, notably the clever initialization schemes and the automatic "unconstraining" transforms.

In any case, I think I have my answer (I can't do what I want) but perhaps what I'm doing above is only possible because of these enhancements, so thanks for making them! And if this surfaced some bugs in the intended vmap functionality, I'm glad.
> Of course, this particular example is easy to do with plate or auto broadcasting (or even a for loop) but there are more complex cases that are quite tricky and jax.vmap would be helpful
> This works fine, but I do lose some of numpyro's goodies, notably the clever initialization schemes and the automatic "unconstraining" transforms.

I think that it would be useful if I could see an example where using `jax.vmap` instead of, at least, batching-based semantics in `numpyro` increases code simplicity.

> I already have pretty full support for numpyro as a backend with arbitrary uses of vmap in defining models

Interesting! I'd be interested to see this feature in action in some code snippet, if you can point me to one.

> This works fine, but I do lose some of numpyro's goodies, notably the clever initialization schemes and the automatic "unconstraining" transforms.
By the way, that whole `validate_args` class attribute handling is not thread-safe @fehiepsi. We might want to change the mechanism to have it use thread-local variables.

Resolving patch:
diff --git a/numpyro/distributions/distribution.py b/numpyro/distributions/distribution.py
--- a/numpyro/distributions/distribution.py
+++ b/numpyro/distributions/distribution.py
@@ -158,10 +158,10 @@ def gather_pytree_data_fields(cls):
         return all_pytree_data_fields
 
     @classmethod
-    def gather_pytree_aux_fields(cls):
+    def gather_pytree_aux_fields(cls) -> tuple:
         bases = inspect.getmro(cls)
 
-        all_pytree_aux_fields = ()
+        all_pytree_aux_fields = ("_validate_args",)
         for base in bases:
             if issubclass(base, Distribution):
                 all_pytree_aux_fields += base.__dict__.get("pytree_aux_fields", ())
@@ -203,11 +203,15 @@ def tree_unflatten(cls, aux_data, params):
         for k, v in pytree_aux_fields_dict.items():
             setattr(d, k, v)
 
+        # disable args validation during `tree_unflatten` it is called by jax with
+        # placeholder attributes that would make validation fail
+        d._validate_args = False
         Distribution.__init__(
             d,
             pytree_aux_fields_dict["_batch_shape"],
             pytree_aux_fields_dict["_event_shape"],
         )
+        d._validate_args = pytree_aux_fields_dict["_validate_args"]
         return d
 
     @staticmethod


Test patch:
diff --git a/test/test_distributions.py b/test/test_distributions.py
--- a/test/test_distributions.py
+++ b/test/test_distributions.py
@@ -3092,6 +3092,21 @@ def sample(d: dist.Distribution):
         assert samples_batched_dist.shape == (1, *samples_dist.shape)
 
 
+def test_vmap_validate_args():
+    # Test for #1684: vmapping distributions whould work when `validate_args=True`
+    v_dist = jax.vmap(
+        lambda loc, scale: dist.Normal(loc=loc, scale=scale, validate_args=True),
+        in_axes=(0, 0),
+    )(jnp.zeros((2,)), jnp.zeros((2,)))
+
+    # non-regression test
+    v_dist = jax.vmap(
+        lambda loc, scale: dist.Normal(loc=loc, scale=scale, validate_args=False),
+        in_axes=(0, 0),
+    )(jnp.zeros((2,)), jnp.zeros((2,)))
+    assert not v_dist._validate_args
+
+
 def test_multinomial_abstract_total_count():
     probs = jnp.array([0.2, 0.5, 0.3])
     key = random.PRNGKey(0)