Currently there are no untransformed values captured in trace object

[rrkarim]

There is no untransformed values (if distribution requires transformation) in trace object, since initialize_state turns state into sampling_state. Maybe there is a need to capture untransformed values before initialize_state call and then inverting appropriate RVs to get more convenient trace. How one can implement it better?

[rrkarim]

Also why not renaming initialize_state to initialize_sampling_state (also add additional state parameter) to remove redundant code in build_logp_and_deterministic_functions with all the state/observed checks?

[twiecki]

Those sound like reasonable changes and requests, happy to review PRs!

[rrkarim]

Working on it, guess there are some related issues to fix first.

[ColCarroll]

@rrkarim are you making progress on this? i may try to tackle it today. I want to be able to just

data = pm.sample(model())

and have data be the arviz object. This will require:

• Update import paths to get sample at the top level
• Wrap the return value of sample into the arviz object
• Add the untransformed values into the trace (probably as Deterministics? this looks at least a little tricky)

I think the first two will be easy, and I will tackle those first.

[rrkarim]

@ColCarroll, yes, the third one is, indeed, looks little tricky with my knowledge of the current design requirements. So yeah, change what you need to.

[lucianopaz]

I wouldn't add them as Deterministics. The executors are used in different contexts for different purposes. The transformed executor is used in sample and in sample_posterior_predictive, and the latter relies on distributions being in untransformed, transformed or observed values in the sampling state, but never in deterministics. I can tackle this issue later this week but I would do the following:

1. The deterministic values are collected with a callback that does a forward pass through the model, so it also collects the untransformed distribution values. I would simply change this callback to also return the relevant subset of untransformed variables that should be included into the trace.
2. To know the list of distributions that should be stored, we could either rely on the test forward pass that we do in build logp function, or we could add an attribute to the sampling state, that stores the names of the autotransformed distributions.

[ColCarroll]

Thanks, @lucianopaz and @rrkarim!

I think I got to the same place as you are suggesting with 1. I am just finding it hard to figure out where to add that collection logic! :grin:

I am relying on the nice pymc4.utils.NameParts.is_transformed to decide if a distribution should be stored.