Repository: pyro-ppl/numpyro

Issue:
Document behavior of primitive statements in README
Hello! I'll preface this issue by stating that I'm new to numpyro, so there's a significant chance that I'm not using `numpyro.sample` correctly.

When building models, I want to test draws from my random variables to make sure I'm passing parameters correctly, getting the kwargs correct, etc.; I'll lift a line out of a model function, [like this one](https://github.com/pyro-ppl/numpyro/blob/master/examples/baseball.py#L111), and run it in an IPython window to see if I get an error. It seems like I'm unable to do this by itself in an IPython console.

## Minimal example

```
In [1]: import numpyro

In [2]: import numpyro.distributions as d

In [3]: numpyro.__version__
Out[3]: '0.2.0'

In [4]: numpyro.sample("x", d.Normal(0, 1))
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-4-21ce96e72ad0> in <module>
----> 1 numpyro.sample("x", d.Normal(0, 1))

~/miniconda3/envs/fusion/lib/python3.7/site-packages/numpyro/primitives.py in sample(name, fn, obs, sample_shape)
     45     # if there are no active Messengers, we just draw a sample and return it as expected:
     46     if not _PYRO_STACK:
---> 47         return fn(sample_shape=sample_shape)
     48
     49     # Otherwise, we initialize a message...

~/miniconda3/envs/fusion/lib/python3.7/site-packages/numpyro/distributions/distribution.py in __call__(self, *args, **kwargs)
    161
    162     def __call__(self, *args, **kwargs):
--> 163         key = kwargs.pop('random_state')
    164         sample_intermediates = kwargs.pop('sample_intermediates', False)
    165         if sample_intermediates:

KeyError: 'random_state'
```

## Thoughts
I can see that in [numpyro.sample](https://github.com/pyro-ppl/numpyro/blob/master/numpyro/primitives.py#L47), we return `fn(sample_shape=sample_shape)`, which in this case attempts to call `dist.Normal(0, 1)(sample_shape=sample_shape)`. However, looking at [distributions.py](https://github.com/pyro-ppl/numpyro/blob/master/numpyro/distributions/distribution.py#L163), it seems that `Distribution.sample` expects a `random_state` kwarg that isn't getting passed.

When I do the following, everything is fine:

```
In [5]: import jax.random as random

In [6]: d.Normal(0, 1).sample(random.PRNGKey(12))
/Users/chtu8001/miniconda3/envs/fusion/lib/python3.7/site-packages/jax/lib/xla_bridge.py:114: UserWarning: No GPU/TPU found, falling back to CPU.
  warnings.warn('No GPU/TPU found, falling back to CPU.')
Out[6]: DeviceArray(-0.5135873, dtype=float32)

In [7]: d.Normal(0, 1).sample(random.PRNGKey(12), sample_shape=(4, ))
Out[7]:
DeviceArray([-1.31179953, -0.70821768,  0.18907626, -1.09385514],
            dtype=float32)
```

so I expect that we can't use `numpyro.sample` outside a model (although it's not totally clear to me how that is defined, something to do with Messengers and PYRO_STACK). I'm wondering if this is by design and I should just use the second, working example, or if I'm misunderstanding how `numpyro.sample` should be used.

Thanks!


Discussion:
> so I expect that we can't use numpyro.sample outside a model (although it's not totally clear to me how that is defined, something to do with Messengers and PYRO_STACK). 

That is correct. If you are using this `numpyro.sample` statement outside a model, you should instead consider calling the distribution's `sample` method instead. The reason why it works in the context of inference is that the model code is wrapped in a `seed` handler which provides the random generator seed to distribution instances and splits the seed for subsequent `sample` statements. Refer to this Pyro [tutorial](http://pyro.ai/examples/effect_handlers.html) on effect handling, if you are interested in understanding how this happens. To see this, you can call `handlers.seed(numpyro.sample, random.PRNGKey(1))('a', dist.Normal(0., 1.))`, which should work for your case. 

The motivation for this design is that, unlike in Pyro, JAX uses a functional PRNG and mandates passing an explicit rng key, which can get quite complex to pass around in large models, and ensure that it isn't reused by splitting it in downstream `sample` statements. The seed handler takes care of that while ensuring that the user does not have to clutter their model code threading the rng key.


I am closing this issue. Feel free to ask any follow up questions here, or ideally on the [forum](forum.pyro.ai), which will have better discoverability for other users!
> I'm misunderstanding how numpyro.sample should be used.

You already get the main point, that it is by design. We can support `numpyro.sample("x", d.Normal(0, 1), random_state=rng)` or something likes that but it seems redundant to me (using `d.Normal(0, 1).sample(key=rng)` as you did already gives what you need). Personally, I would consider the primitives are meant to be used inside a stochastic function whose random state can be initialized using a random seed (this is the purpose of `seed` handler). If I want to play around with distributions themselves, then I won't use `sample` primitives. However, we will keep this in mind and if it turns out that this design causes much trouble to users, we will rethink about it.

Welcome to NumPyro!
Thank you both for your replies. That makes sense and I appreciate the clarification!

> However, we will keep this in mind and if it turns out that this design causes much trouble to users, we will rethink about it.

I don't think the design is problematic at all. But it may be helpful for this to be documented in [numpyro.sample](http://num.pyro.ai/en/stable/primitives.html#numpyro.primitives.sample), since it is different from `pyro.sample`. The Pyro tutorial [includes](http://pyro.ai/examples/intro_part_i.html#The-pyro.sample-Primitive) running `pyro.sample` as a standalone command, so it may be surprising to others that this won't work in NumPyro.
> But it may be helpful for this to be documented in numpyro.sample,

Good point! We should document it. Thanks!! :)

Resolving patch:
diff --git a/numpyro/primitives.py b/numpyro/primitives.py
--- a/numpyro/primitives.py
+++ b/numpyro/primitives.py
@@ -60,21 +60,28 @@ def __call__(self, *args, **kwargs):
             return self.fn(*args, **kwargs)
 
 
-def sample(name, fn, obs=None, sample_shape=()):
+def sample(name, fn, obs=None, random_state=None, sample_shape=()):
     """
     Returns a random sample from the stochastic function `fn`. This can have
     additional side effects when wrapped inside effect handlers like
     :class:`~numpyro.handlers.substitute`.
 
+    .. note::
+        By design, `sample` primitive is meant to be used inside a NumPyro model.
+        Then :class:`~numpyro.handlers.seed` handler is used to inject a random
+        state to `fn`. In those situations, `random_state` keyword will take no
+        effect.
+
     :param str name: name of the sample site
     :param fn: Python callable
     :param numpy.ndarray obs: observed value
+    :param jax.random.PRNGKey random_state: an optional random key for `fn`.
     :param sample_shape: Shape of samples to be drawn.
     :return: sample from the stochastic `fn`.
     """
     # if there are no active Messengers, we just draw a sample and return it as expected:
     if not _PYRO_STACK:
-        return fn(sample_shape=sample_shape)
+        return fn(random_state=random_state, sample_shape=sample_shape)
 
     # Otherwise, we initialize a message...
     initial_msg = {
@@ -82,7 +89,7 @@ def sample(name, fn, obs=None, sample_shape=()):
         'name': name,
         'fn': fn,
         'args': (),
-        'kwargs': {'sample_shape': sample_shape},
+        'kwargs': {'random_state': random_state, 'sample_shape': sample_shape},
         'value': obs,
         'scale': 1.0,
         'is_observed': obs is not None,


Test patch: