Numpyro: Document behavior of primitive statements in README

[tuchandra]

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, 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, 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, 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!

[neerajprad]

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 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.

[neerajprad]

I am closing this issue. Feel free to ask any follow up questions here, or ideally on the forum, which will have better discoverability for other users!

[fehiepsi]

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!

[tuchandra]

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, since it is different from pyro.sample. The Pyro tutorial includes running pyro.sample as a standalone command, so it may be surprising to others that this won't work in NumPyro.

[fehiepsi]

But it may be helpful for this to be documented in numpyro.sample,

Good point! We should document it. Thanks!! :)

[timnon]

Ran into the same problem, also expected a similar behaviour as in pyro. It would also be nice to have the option to sample from a complete model for testing just by running the model-function without any modifications.

[neerajprad]

@timnon - Thanks for the feedback. We have tried to keep the interface as similar as possible, but there are a few things that cannot be done in a purely functional land. The particular issue in this case is that sample statements will need an explicit PRNGKey to generate samples, since there is no global random state.

Instead of calling model(), this requires that users call seed(model, random.PRNGKey(1))() in NumPyro. We found this to be the best compromise between keeping the interface the same as Pyro vs. having users use explicit random number generator keys in sample statements and using random.split in their models. We will be happy to consider alternate solutions to this issues, so please let us know if you have any suggestions.

[neerajprad]

I have reopened this issue so that this is better highlighted in the README.

[timnon]

Thanks, this seems to be a deeply technical issue, but solution with the seed-function works well. For people looking for some minimal copy-paste code how to sample from some numpyro model:

import numpyro
import jax
import random

def model():
    return numpyro.sample("obs", numpyro.distributions.Normal(0, 1.))

some_random_int = random.randint(0,10000)
sample = numpyro.handlers.seed(model,jax.random.PRNGKey(some_random_int))()
print(sample)

[timnon]

P.S. the correct way to sample with jax is probably:

import numpyro
import jax

def model():
    return numpyro.sample("obs", numpyro.distributions.Normal(0, 1.))

n_samples = 100
samples = list()
key = jax.random.PRNGKey(0)
for _ in range(n_samples):
    key,key_ = jax.random.split(key)
    sample = float(numpyro.handlers.seed(model,key)())
    samples.append(sample)    
print(samples)

[neerajprad]

@timnon - You can also use jax.vmap to vectorize the sampling (instead of generating samples sequentially):

update: FIXED syntax below

rngs = random.split(random.PRNGKey(0), n_samples)
samples = vmap(lambda key: numpyro.handlers.seed(model, key)())(rngs)

[fehiepsi]

@timnon probably you want numpyro.handlers.seed(model,key_) instead, because key is split in the next iteration.

The reason why we don't use numpy global random state is those states will be constant under jax transformations such as jit, vmap. That means if we use global random state, we can get a sample from model by calling model(), and to get another sample, we can keep calling model(). However, when we jit model jmodel = jit(model), then sample1 = jmodel(), sample2 = jmodel() will give constant samples (sample1 is equal to sample2). So although using seed is a bit inconvenient, it is a safe way to avoid that mistake.

You can also use jax.vmap to vectorize the sampling

Yeah, I intended to add sample_shape/num_samples arg to predictive utility so that users can get prior samples without having to use return or trace. :D

One thing I also found very useful is to allow using predictive to get returned values. Those returned values need not to be part of the probabilistic model, but are transformed values of some latent variables. In Stan, they are parameters in generated_quantities. In PyMC3, they are pm.Deterministic nodes. I guess in Pyro, it is stored in _RETURN node. I think we can follow Pyro approach for this feature. WDYT?

[timnon]

@neerajprad I tried the vectorized approach (i think the rngs should be a parameter of the vectorised method):

import numpyro
import jax

def model():
    return numpyro.sample("obs", numpyro.distributions.Normal(0, 1.))

n_samples = 100
rngs = jax.random.split(jax.random.PRNGKey(0), n_samples)
samples = jax.vmap(lambda key: numpyro.handlers.seed(model, key))(rngs)
print(samples)

However, i get TypeError: <class 'numpyro.handlers.seed'> is not a valid Jax type . I never used jax before, but i think most people will first try to sample from some model when starting with numpyro for some sniff tests.

[fehiepsi]

@timnon It is a typo: numpyro.handlers.seed(model, key) -> numpyro.handlers.seed(model, key)(). numpyro.handlers.seed(model, key) is just "another" model with seed effect, we need to call it to get returned values.

i think most people will first try to sample from some model when starting with numpyro for some sniff tests

Agreed! I will make a predictive utility for this purpose today, so given a model, you can simply do predictive(rng, model, num_samples) to get num_samples samples from all latent sites of your model. :)

[neerajprad]

As @fehiepsi mentioned, that was a mistake in the syntax; fixed above.

[timnon]

thanks @fehiepsi . For the sake of completeness, here is the complete best way to sample:

import numpyro
import numpyro.handlers
import jax

def model():
    return numpyro.sample("obs", numpyro.distributions.Normal(0, 1.))

def sample_model(model,n_samples=100):
    keys = jax.random.split(jax.random.PRNGKey(0), n_samples)
    return jax.vmap(lambda key: numpyro.handlers.seed(model, key)())(keys)
print(sample_model(model))

wouldnt it be easy to include some sample_model-method?

[fehiepsi]

wouldnt it be easy to include some sample_model-method?

@timnon Yeah, that is simple to add (we will add it soon). We have mostly focused on building inference engines. Given that these engines are stable now, we should focus more on user experiences. Feedbacks from you and @tuchandra in this topic are very valuable to us, thanks a lot!

[tuchandra]

I'm glad it was useful :) thank you for being so responsive and helpful!