# Path: blackjax/types.py



from functools import partial
from typing import Callable
from blackjax.types import PRNGKey
import jax
import jax.numpy as jnp
""" All things resampling. """




def _resampling_func(func, name, desc="", additional_params="") -> Callable:
    # Decorator for resampling function

    doc = f"""
    {name} resampling. {desc}

    Parameters
    ----------
    weights: jnp.ndarray
        Weights to resample
    key: jnp.ndarray
        PRNGKey to use in resampling
    {additional_params}

    Returns
    -------
    idx: jnp.ndarray
        Array of integers to use for resampling
    """

    func.__doc__ = doc
    return func


@partial(_resampling_func, name="Systematic")

# Path: blackjax/types.py



# Path: blackjax/smc/resampling.py
from functools import partial
from typing import Callable
from blackjax.types import PRNGKey
import jax
import jax.numpy as jnp

""" All things resampling. """




def _resampling_func(func, name, desc="", additional_params="") -> Callable:
    # Decorator for resampling function

    doc = f"""
    {name} resampling. {desc}

    Parameters
    ----------
    weights: jnp.ndarray
        Weights to resample
    key: jnp.ndarray
        PRNGKey to use in resampling
    {additional_params}

    Returns
    -------
    idx: jnp.ndarray
        Array of integers to use for resampling
    """

    func.__doc__ = doc
    return func


@partial(_resampling_func, name="Systematic")

def systematic(weights: jnp.ndarray, rng_key: PRNGKey) -> jnp.ndarray: