# Path: blackjax/types.py



from typing import Callable
from blackjax.types import PyTree
import jax
import jax.numpy as jnp


def ess(log_weights: jnp.ndarray, log: bool = True) -> float:
    """Compute the effective sample size from the log-weights in a numerically stable way.

    Parameters
    ----------
    log_weights: np.ndarray
        log-weights of the sample
    log: bool
        Compute the log-ESS or the ESS

    Returns
    -------
    ess: float
        The effective sample size
    """

    log_weights = log_weights - jnp.max(log_weights)
    w = jnp.exp(log_weights)
    if log:
        w2: jnp.ndarray = jnp.exp(2 * log_weights)
        res = 2 * jnp.log(w.sum()) - jnp.log(w2.sum())
    else:
        res = jnp.sum(w) ** 2 / jnp.sum(w ** 2)
    return res


def ess_solver(
    logprob_fn: Callable,

# Path: blackjax/types.py



# Path: blackjax/smc/ess.py
from typing import Callable
from blackjax.types import PyTree
import jax
import jax.numpy as jnp



def ess(log_weights: jnp.ndarray, log: bool = True) -> float:
    """Compute the effective sample size from the log-weights in a numerically stable way.

    Parameters
    ----------
    log_weights: np.ndarray
        log-weights of the sample
    log: bool
        Compute the log-ESS or the ESS

    Returns
    -------
    ess: float
        The effective sample size
    """

    log_weights = log_weights - jnp.max(log_weights)
    w = jnp.exp(log_weights)
    if log:
        w2: jnp.ndarray = jnp.exp(2 * log_weights)
        res = 2 * jnp.log(w.sum()) - jnp.log(w2.sum())
    else:
        res = jnp.sum(w) ** 2 / jnp.sum(w ** 2)
    return res


def ess_solver(
    logprob_fn: Callable,

    particles: PyTree,