def _autocorr(x, axis=-1, center=True):
        """Compute auto-correlation along one axis."""

        dtype = x.dtype
        shift = (-1 - axis) if axis < 0 else (len(x.shape) - 1 - axis)
        x = jnp.transpose(x, jnp.roll(jnp.arange(len(x.shape)), shift))
        if center:
            x -= x.mean(axis=-1, keepdims=True)

        # Zero pad to the next power of 2 greater than 2 * x_len
        x_len = x.shape[-1]
        pad_len = int(2.0 ** jnp.ceil(jnp.log2(x_len * 2)) - x_len)
        x = jnp.pad(x, (0, pad_len))[:-pad_len]

        # Autocorrelation is IFFT of power-spectral density
        fft = jnp.fft.fft(x.astype(jnp.complex64))
        prod = jnp.fft.ifft(fft * jnp.conj(fft))
        prod = jnp.real(prod[..., :x_len]).astype(dtype)

        # Divide to obtain an unbiased estimate of the expectation
        denominator = x_len - jnp.arange(0.0, x_len)
        res = prod / denominator
        return jnp.transpose(res, jnp.roll(jnp.arange(len(res.shape)), -shift))
Compute auto-correlation along one axis.
fortuna/prob_model/posterior/sgmcmc/sgmcmc_diagnostic.py