Memory mapping `arviz.InferenceData`

[OriolAbril]

@Benjamin

shape=(chain, draw, #of RVs in the model) and that’s why it’s effectively doubling the memory footprint?

Equivalent to this 2nd one, but instead of a single array with this shape you have multiple arrays. In your case,

@Benjamin

    coords = {"group": frame["group"].unique()}
    with pm.Model(coords=coords) as model:
        # Data
        g = pm.MutableData("g", frame["group"])
        x = pm.MutableData("x", frame["x"])
        y = pm.ConstantData("y", frame["y"])

        # Population level
        alpha = pm.Normal("alpha")
        beta = pm.Normal("beta")

        omega = pm.HalfNormal("omega")
        sigma = pm.HalfNormal("sigma")
        tau = pm.HalfNormal("tau")

        # Group level
        epsilon_alpha_g = pm.ZeroSumNormal("epsilon_alpha_g", dims="group")
        alpha_g = pm.Deterministic("alpha_g", alpha + omega * epsilon_alpha_g)

        epsilon_beta_g = pm.ZeroSumNormal("epsilon_beta_g", dims="group")
        beta_g = pm.Deterministic("beta_g", beta + tau * epsilon_beta_g)

        # Linear model
        mu_ = pm.Deterministic("mu_", alpha_g[g] + beta_g[g] * x)

        # Likelihood
        pm.Normal("y_obs", mu=mu_, sigma=sigma, observed=y, shape=mu_.shape)

you have alpha, beta, omega, sigma, tau with shape chain, draw, epsilon_alpha_g, alpha_g, epsilon_beta_g, beta_b with shape chain, draw, group and mu with chain, draw, n_observations. Thus, your posterior is equivalent to an array of floats with shape chain, draw, 5 + n_group * 4 + n_observations. The log prior will not have the deterministics, but for the variables it has, they’ll be the same shape as in the posterior. Therefore in your case (unless dropping all deterministics) it wouldn’t quite double the memory needed, it would be equivalent to an array of shape chain, draw, 5 + 2 * n_groups.

You only have a single observed variable y_obs, so your log_likelihood and posterior_predictive will be a single array with shape chain, draw, n_observations and you’d be in the case where the posterior needs more memory than the log_likelihood group. Once you stop storing mu this is inverted though.