code: def _ess_ufunc(ary):
      
      target = np.empty(ary.shape[:-2])
      for idx in np.ndindex(target.shape):
          target[idx] = _get_ess(ary[idx])
      return target
docs: 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'),))