# Path: group_testing/utils.py
# def unique(rng, binary_vectors: np.ndarray) -> int:
# def select_from_sizes(values: np.ndarray, sizes: np.ndarray) -> np.ndarray:
# 
# Path: group_testing/samplers/sampler.py
# class Sampler:
#   def __init__(self):
#   def reset(self):
#   def is_cheap(self):
#   def marginal(self):
#   def reset_convergence_metric(self):


from absl import logging
from group_testing import utils
from group_testing.samplers import sampler
import gin
import jax
import jax.nn
import jax.numpy as np
import jax.scipy.special
    return np.stack((alpha, beta), axis=0)

  def cond_fun(alpha_beta_iteration):
    alphabeta, iteration = alpha_beta_iteration
    marginal = marginal_from_alphabeta(alphabeta)
    marginal_plus_one_iteration = marginal_from_alphabeta(
        lbp_loop(0, alphabeta))
    converged = np.allclose(marginal, marginal_plus_one_iteration, atol=atol)
    return (not converged) and (iteration < max_iterations)

  def body_fun(alpha_beta_iteration):
    alphabeta, iteration = alpha_beta_iteration
    alphabeta = jax.lax.fori_loop(0, min_iterations, lbp_loop, alphabeta)
    iteration += min_iterations
    return [alphabeta, iteration]

  # Run LBP while loop
  while cond_fun(alpha_beta_iteration):
    alpha_beta_iteration = body_fun(alpha_beta_iteration)

  alphabeta, _ = alpha_beta_iteration

  # Compute two consecutive marginals
  marginal = marginal_from_alphabeta(alphabeta)
  marginal_plus_one_iteration = marginal_from_alphabeta(lbp_loop(0, alphabeta))

  return marginal, np.amax(np.abs(marginal - marginal_plus_one_iteration))


@gin.configurable

# Path: group_testing/utils.py
# def unique(rng, binary_vectors: np.ndarray) -> int:
# def select_from_sizes(values: np.ndarray, sizes: np.ndarray) -> np.ndarray:
# 
# Path: group_testing/samplers/sampler.py
# class Sampler:
#   def __init__(self):
#   def reset(self):
#   def is_cheap(self):
#   def marginal(self):
#   def reset_convergence_metric(self):


# Path: group_testing/samplers/loopy_belief_propagation.py
from absl import logging
from group_testing import utils
from group_testing.samplers import sampler
import gin
import jax
import jax.nn
import jax.numpy as np
import jax.scipy.special

    return np.stack((alpha, beta), axis=0)

  def cond_fun(alpha_beta_iteration):
    alphabeta, iteration = alpha_beta_iteration
    marginal = marginal_from_alphabeta(alphabeta)
    marginal_plus_one_iteration = marginal_from_alphabeta(
        lbp_loop(0, alphabeta))
    converged = np.allclose(marginal, marginal_plus_one_iteration, atol=atol)
    return (not converged) and (iteration < max_iterations)

  def body_fun(alpha_beta_iteration):
    alphabeta, iteration = alpha_beta_iteration
    alphabeta = jax.lax.fori_loop(0, min_iterations, lbp_loop, alphabeta)
    iteration += min_iterations
    return [alphabeta, iteration]

  # Run LBP while loop
  while cond_fun(alpha_beta_iteration):
    alpha_beta_iteration = body_fun(alpha_beta_iteration)

  alphabeta, _ = alpha_beta_iteration

  # Compute two consecutive marginals
  marginal = marginal_from_alphabeta(alphabeta)
  marginal_plus_one_iteration = marginal_from_alphabeta(lbp_loop(0, alphabeta))

  return marginal, np.amax(np.abs(marginal - marginal_plus_one_iteration))


@gin.configurable

class LbpSampler(sampler.Sampler):