def _ess_ufunc(ary):
    """Ufunc for computing effective sample size.

    This can be used on an xarray Dataset, using
    `xr.apply_ufunc(_ess_ufunc, ..., input_core_dims=(('chain', 'draw'),))
    """
    target = np.empty(ary.shape[:-2])
    for idx in np.ndindex(target.shape):
        target[idx] = _get_ess(ary[idx])
    return target
Ufunc for computing effective sample size.

    This can be used on an xarray Dataset, using
    `xr.apply_ufunc(_ess_ufunc, ..., input_core_dims=(('chain', 'draw'),))