Substituting a plate takes range?

[OlaRonning]

Using the handler.substitute with a plate is counter-intuitive. Given that plate takes an int for the size parameter, I would expect the substitute value to also have type int. However, the substitution need a range [0,n). Below is a small example that illustrates this behavior:

from jax import numpy as jnp
from numpyro import plate, sample, prng_key
from numpyro.handlers import seed, substitute
from numpyro.distributions import Normal

def model(y):
    n = y.shape[0] if y is not None else 1
    with plate('n', n):
        y = sample('y', Normal(), obs=y)
    return y

with seed(rng_seed=0), substitute(data={'n': jnp.arange(10)}):  # Wrong type
    y = model(None)

assert jnp.unique(y).shape == (10,)

with seed(rng_seed=0), substitute(data={'n': 10}):
    y = model(None)  # <-- **AttributeError: 'int' object has no attribute 'shape'**

[fehiepsi]

The plate value is an index array. Similar for subsample. If you want to change size, you can provide it as an argument or a global configuration.

[OlaRonning]

I see. Thanks for clarifying the rational.