radon notebook cell 12 broken

[MooersLab]

model = hierarchical_model_nc(data, county_idx) az_trace_nc = sample(model)

...site-packages/tensorflow_probability/python/mcmc/nuts.py:413 bootstrap_results len(init_state), len(step_size)))

ValueError: Expected either one step size or 7 (size of `init_state`), but found 9

[mstump]

The radon notebook has since further degraded. It wasn't updated after the return type for sample was changed to InterfaceData. I've started an attempt to clean it up but I'm not entirely sure what the previous return type was, I just know it loosely conformed to a Dict type. Perhaps it was an xarray? As part of my attempt would the team be open to me adding Python 3 type annotations? I find that it really helps in understanding an API and can catch interface drift. I know some people feel strongly about the subject so I wanted to check first.

      3     pm4_trace, _ = pm.inference.sampling.sample(
----> 4         model, num_chains=init_num_chains, num_samples=10, burn_in=10, step_size=1., xla=True)
      5     for i in range(3):
      6         step_size_ = []

TypeError: cannot unpack non-iterable InferenceData object

[aloctavodia]

Hi, @mstump we are already using type annotations so any improvement in that are will be very welcome. The previous type was a dict, if you want to recover the old behaviour you could comment this line https://github.com/pymc-devs/pymc4/blob/master/pymc4/inference/sampling.py#L160 and instead use:

return posterior, sampler_stats

The code below should fix the arviz integration part, but the error reported by @MooersLab should still there. May be @junpenglao has a better idea of what is going on.

def sample(model, init_num_chains=50, num_samples=500, burn_in=500):
    init_num_chains = 50
    trace = pm.sample(model,
                      num_chains=init_num_chains,
                      num_samples=10,
                      burn_in=10,
                      step_size=1.,
                      xla=True).posterior

    for i in range(3):
        step_size_ = []
        for _, x in trace.items():
            std = tf.math.reduce_std(x, axis=[1,0])
            step_size_.append(
                std[tf.newaxis, ...] * tf.ones([init_num_chains] + std.shape,
                                               dtype=std.dtype))

        trace = pm.sample(model,
                          num_chains=init_num_chains,
                          num_samples=10 + 10*i,
                          burn_in=10 + 10*i,
                          step_size=step_size_,
                          xla=True).posterior

    num_chains = 5
    step_size_ = []
    for _, x in trace.items():
        std = tf.math.reduce_std(x, axis=[1, 0])
        step_size_.append(std[tf.newaxis, ...] * tf.ones([num_chains]+std.shape,
                                                         dtype=std.dtype))

    trace = pm.sample(model, num_chains=num_chains, num_samples=num_samples,
                      burn_in=burn_in, step_size=step_size_, xla=True)

    return trace

[lucianopaz]

I can't check it out now, but the error seems to indicate there are some extra parameters that are being included in the initial value for the step. 9 instead of 7. I think this could mean that for some reason, an auto transformed parameter is leaking into the initial state.

[junpenglao]

exact reason of what @lucianopaz said. The chain now contains also the transformed variable (in Real), thus messing up the mapping to step size.

[mstump]

@junpenglao that's only sort of correct, I went back in time trying different SHAs to see when the last time this notebook worked, it was 28e86b7. In that version of the notebook the keys of the trace are below and it includes a couple transformed variables, but no eps:

hierarchical_model/mu_alpha
hierarchical_model/__log_sigma_alpha
hierarchical_model/mu_beta
hierarchical_model/alpha
hierarchical_model/beta
hierarchical_model/__log_sigma_beta 
hierarchical_model/__log_eps

The collab notebook is here

In master I restricted the loop the untransformed variables getting the list of keys from the evaluated model state:

    _, evaluated_state = pm.evaluate_model(model)   
    untransformed_values = evaluated_state.untransformed_values.keys()
    print(untransformed_values)
    print(len(untransformed_values))

Which gives me the set:

dict_keys(['hierarchical_model/mu_alpha', 'hierarchical_model/sigma_alpha', 'hierarchical_model/mu_beta', 'hierarchical_model/sigma_beta', 'hierarchical_model/alpha', 'hierarchical_model/beta', 'hierarchical_model/eps'])

After modifying the cell to the below I got it to run:

def sample(model, init_num_chains=50, num_samples=500, burn_in=500):
    _, evaluated_state = pm.evaluate_model(model)   
    untransformed_values = evaluated_state.untransformed_values.keys()

    init_num_chains = 50
    trace = pm.sample(model,
                      num_chains=init_num_chains,
                      num_samples=10,
                      burn_in=10,
                      step_size=1.,
                      xla=True).posterior

    for i in range(3):
        step_size_ = []
        for variable_name in trace.keys():
            if variable_name not in untransformed_values:
                continue

            x = trace[variable_name]
            std = tf.math.reduce_std(x, axis=[1,0])
            step_size_.append(
                std[tf.newaxis, ...] * tf.ones([init_num_chains] + std.shape,
                                               dtype=std.dtype))

        trace = pm.sample(model,
                          num_chains=init_num_chains,
                          num_samples=10 + 10*i,
                          burn_in=10 + 10*i,
                          step_size=step_size_,
                          xla=True).posterior

    num_chains = 5
    step_size_ = []
    for variable_name in trace.keys():
        if variable_name not in untransformed_values:
            continue

        x = trace[variable_name]
        std = tf.math.reduce_std(x, axis=[1, 0])
        step_size_.append(std[tf.newaxis, ...] * tf.ones([num_chains]+std.shape,
                                                         dtype=std.dtype))

    trace = pm.sample(model, num_chains=num_chains, num_samples=num_samples,
                      burn_in=burn_in, step_size=step_size_, xla=True)

    return trace

I'm not sure if the exclusion of the transformed variables that were present in the historical commit have a significant impact. I'm coming from an engineering background so some of this is still hand waving for me. I'm not entirely sure why we're calculating the step size rather than relying on the implementations already present in TF probability, they have DualAveragingStepSizeAdaptation, is this code redundant?

hierarchical_model/__log_sigma_alpha
hierarchical_model/__log_sigma_beta 
hierarchical_model/__log_eps

Also, it feels like the list of transformed and untransformed variable names should be added to metadata of model as attributes (I'm not entirely sure if this is possible). Adding it to the output of sample doesn't feel right because it's an an InterfaceData which is defined in another package and we'd have to change the interface.

[disadone]

I follow the code above,

the line

std = tf.math.reduce_std(x, axis=[1,0])

gives

Exception has occurred: AttributeError 'numpy.dtype' object has no attribute 'is_complex'

using vscode remote container