Add deterministic primitive

[fehiepsi]

This PR adds deterministic primitive to Pyro to record deterministic values in a trace as requested in pyro-ppl/numpyro#462, #2063 (comment), and here and there in the forum. The functionality is similar to pymc3.Deterministic or Stan's transformed parameters and generated quantities.

Note that the implementation is just an alias of sample(name, dist.Delta(value), obs=value) so it might have some overhead to computing joint density. However, I think that this primitive is mostly used in prediction, so users can remove this overhead by using the pattern

if y is None:
    pyro.deterministic(...)

Edit: It seems that we also need an additional arg event_dim to explicitly state the event dimension of that deterministic site (see the test) so that log_prob shape matches those of the other sites. This seems like a bad idea because we only need to record the values of these deterministic sites. Maybe we should require that these deterministic sites should not be used inside a plate statement. In case it is used inside a plate state, an explicit event_dim is required.

[fritzo]

Note dist.Delta(...).mask(False) may remove the tensor op overhead, as e.g. exploited in #2110, #2124

[fehiepsi]

Thanks, that would be great! I'll try to use it.

[neerajprad]

This looks great, pending any further review comments from @fritzo.

[neerajprad]

Maybe we should require that these deterministic sites should not be used inside a plate statement. In case it is used inside a plate state, an explicit event_dim is required.

I think your current setup is fine, but this will be a reasonable restriction to have. I am not sure if there are any use cases for deterministic inside a plate context, but it may be easier to just have it than to reliably disable it.

[fritzo]

Ha ha this reminds me of pyro.managed #237, which @eb8680 found very ugly. @eb8680 do you think a new pyro.deterministic deserves a new site type?

[eb8680]

do you think a new pyro.deterministic deserves a new site type?

Sure, this seems like a simple way to satisfy the intended use case (predicting deterministic quantities in Predictive).

[eb8680]

Could you expand the docstring with information about what exactly is added to the trace and when someone might want to use deterministic?

[neerajprad]

In case you find it useful, we needed this for the linear regression example to record the mean response, which will be a deterministic function of the coefficients, data and the intercept.

[fritzo]

@fehiepsi you might want to add to the docstring an EXPERIMENTAL marker and a comment that "this currently converts to a pyro.sample statement, but may change in the future", then we can later create a new site type and remove the event_dim kwarg.

[fritzo]

@fehiepsi do you want to make the docs changes? I've been holding off merging because I interpreted your 👍 as "sure I'll make those changes before this merges". But if you want to make those changes in a follow-up PR that is also fine.

[fehiepsi]

@fritzo I want to make the docs changes. I have marked "awaiting response" to remind me that I should do it. :)

[fehiepsi]

@eb8680 @neerajprad @fritzo Sorry for the delay and thanks for your suggestions!

[fritzo]

@eb8680 @neerajprad do you have further comments before we merge?

[neerajprad]

LGTM.

[eb8680]

LGTM