// make sure to handle scalars correctly - add extra dim if needed
if len(x.shape) == 2:
is_scalar = True
x = np.atleast_3d(mtrace.get_values(var, combine=False))
else:
is_scalar = False
// now we are going to transpose all dims - makes the loop below
// easier by moving the axes of the variable to the front instead
// of the chain and sample axes
x = x.transpose()
Vhat = get_vhat(x)
// get an array the same shape as the var
_n_eff = np.zeros(x.shape[:-2])
// iterate over tuples of indices of the shape of var
for tup in np.ndindex(*list(x.shape[:-2])):
_n_eff[tup] = get_neff(x[tup], Vhat[tup])
// we could be using np.squeeze here, but we don"t want to squeeze
// out dummy dimensions that a user inputs
if is_scalar:
n_eff[var] = _n_eff[0]
else:
// make sure to transpose the dims back
n_eff[var] = np.transpose(_n_eff)
return n_eff
After Change
if not isinstance(mtrace, MultiTrace):
// Return neff for non-multitrace array
return generate_neff(mtrace)
if mtrace.nchains < 2:
raise ValueError(
"Calculation of effective sample size requires multiple chains "
