def kernel_stein_discrepancy_imq(
    samples: List[PyTree],
    grads: List[PyTree],
    c: float = 1.0,
    beta: float = -0.5,
) -> float:
    """Kernel Stein Discrepancy with the Inverse Multiquadric (IMQ) kernel.

    See `Gorham J. and Mackey L., 2017 <https://proceedings.mlr.press/v70/gorham17a/gorham17a.pdf>`_ for more details.

    Parameters
    ----------
        samples: List[PyTree]
            The list of `PyTree`, each representing an MCMC sample.
        grads: List[PyTree]
            The list of the corresponding density gradients.
        c: float
            :math:`c > 0` kernel bias hyperparameter.
        beta: float
            :math:`beta < 0` kernel exponent hyperparameter.

    Returns
    -------
        ksd_img: float
            The kernel Stein discrepancy value.
    """
    if not c > 0:
        raise ValueError("`c` should be > 0.")
    if not beta < 0:
        raise ValueError("`beta` should be < 0.")

    samples = ravel_pytree(samples)[0].reshape(len(samples), -1)
    grads = ravel_pytree(grads)[0].reshape(len(grads), -1)

    def _k_0(param1, param2, grad1, grad2, c, beta):
        dim = param1.shape[0]
        diff = param1 - param2
        base = c**2 + jnp.dot(diff, diff)
        kern = jnp.dot(grad1, grad2) * base**beta
        kern += -2 * beta * jnp.dot(grad1, diff) * base ** (beta - 1)
        kern += 2 * beta * jnp.dot(grad2, diff) * base ** (beta - 1)
        kern += -2 * dim * beta * (base ** (beta - 1))
        kern += -4 * beta * (beta - 1) * base ** (beta - 2) * jnp.sum(jnp.square(diff))
        return kern

    _batched_k_0 = vmap(_k_0, in_axes=(None, 0, None, 0, None, None))

    def _ksd(accum, x):
        sample1, grad1 = x
        accum += jnp.sum(_batched_k_0(sample1, samples, grad1, grads, c, beta))
        return accum, None

    ksd_sum, _ = lax.scan(_ksd, 0.0, (samples, grads))
    return jnp.sqrt(ksd_sum) / samples.shape[0]
Kernel Stein Discrepancy with the Inverse Multiquadric (IMQ) kernel.

    See `Gorham J. and Mackey L., 2017 <https://proceedings.mlr.press/v70/gorham17a/gorham17a.pdf>`_ for more details.

    Parameters
    ----------
        samples: List[PyTree]
            The list of `PyTree`, each representing an MCMC sample.
        grads: List[PyTree]
            The list of the corresponding density gradients.
        c: float
            :math:`c > 0` kernel bias hyperparameter.
        beta: float
            :math:`beta < 0` kernel exponent hyperparameter.

    Returns
    -------
        ksd_img: float
            The kernel Stein discrepancy value.
    
fortuna/prob_model/posterior/sgmcmc/sgmcmc_diagnostic.py