amifalk

import jax
import jax.random as random
from numpyro.infer import SVI

def batch_svi_setup(model, guide, optim, loss, 
                    model_args,
                    *,
                    rng_key,
                    svi_init_params=None,
                    n_init=1,

amifalk

from functools import partial
from operator import attrgetter

import jax
import jax.numpy as jnp
import jax.random as random

@partial(jax.jit, static_argnames=['field_names'])
def collect_fields(state, field_names: tuple):
    """