Sample_shape vs plate

What is the difference between sample_shape and plate? Both seem to accomplish the same thing (getting iid draws). However, I recently switched from use trace_elbo to trace_graph_elbo and now i get a bunch of errors like this:

/opt/jax-source/jax/_src/linear_util.py:190: UserWarning: Missing a plate statement for batch dimension -1 at site 'beta'. You can use `numpyro.util.format_shapes` utility to check shapes at all sites of your model.
  ans = self.f(*args, **dict(self.params, **kwargs))

Things still seem to fit without a problem (I think). Any help here would be greatly appreciated!

1 Like

We encourage users to use plate, so that more sotisphicated inference algorithms can be applied to the model.

Hmm still a bit confused by this response. is the functionality of the two still the same?

plate does more things: it annotates batch dimensions with the names. Some mechanisms, like enumeration (Example: Enumerate Hidden Markov Model β€” NumPyro documentation), require the β€œname” information. For TraceGraph_ELBO, we use name information to check whether the model and guide are compatible, as below. If you think that they are compatible, you can ignore the warning.

    1. Each sample site in the model also appears in the guide and is not
        marked auxiliary.
    2. Each sample site in the guide either appears in the model or is marked,
        auxiliary via ``infer={'is_auxiliary': True}``.
    3. Each :class:`~numpyro.primitives.plate` statement in the guide also
        appears in the model.
    4. At each sample site that appears in both the model and guide, the model
        and guide agree on sample shape.