Index

A

activation (jaxili.compressor.CNN2DCompressor attribute) (jaxili.compressor.MLPCompressor attribute) (jaxili.model.AffineCoupling attribute) (jaxili.model.ConditionalMADE attribute) (jaxili.model.ConditionalMAF attribute) (jaxili.model.ConditionalRealNVP attribute) (jaxili.model.MAFLayer attribute) (jaxili.model.MixtureDensityNetwork attribute)

AffineCoupling (class in jaxili.model) append_simulations() (jaxili.inference.nle.NLE method) (jaxili.inference.npe.NPE method)

B

backward() (jaxili.model.ConditionalMAF method) (jaxili.model.MAFLayer method) batch_stats (jaxili.train.TrainState attribute)

bias (jaxili.model.MaskedLinear attribute) bind_model() (jaxili.train.TrainerModule method) build_posterior() (jaxili.inference.nle.NLE method) (jaxili.inference.npe.NPE method)

C

check_density_estimator() (in module jaxili.utils) check_hparams_maf() (in module jaxili.utils) check_hparams_mdn() (in module jaxili.utils) check_hparams_realnvp() (in module jaxili.utils) CNN2DCompressor (class in jaxili.compressor) compress() (jaxili.model.Compressor_w_NDE method) compressor (jaxili.model.NDE_Compressor attribute) compressor_hparams (jaxili.model.NDE_Compressor attribute)

Compressor_w_NDE (class in jaxili.model) ConditionalMADE (class in jaxili.model) ConditionalMAF (class in jaxili.model) ConditionalRealNVP (class in jaxili.model) create_data_loader() (in module jaxili.utils) create_functions() (jaxili.train.TrainerModule method) create_jitted_functions() (jaxili.train.TrainerModule method) create_trainer() (jaxili.inference.nle.NLE method) (jaxili.inference.npe.NPE method)

D

DirectPosterior (class in jaxili.posterior.direct_posterior)

E

embedding() (jaxili.model.NDE_w_Standardization method)

embedding_net (jaxili.model.NDE_w_Standardization attribute) eval_model() (jaxili.train.TrainerModule method)

F

forward() (jaxili.model.MAFLayer method)

G

gaussian (jaxili.model.ConditionalMADE attribute)

gaussian_kernel_matrix() (in module jaxili.loss) generate_config() (jaxili.train.TrainerModule method)

H

handle_non_serializable() (in module jaxili.utils) hidden_dims (jaxili.model.ConditionalMADE attribute) (jaxili.model.MAFLayer attribute)

hidden_size (jaxili.compressor.MLPCompressor attribute)

I

Identity (class in jaxili.compressor) init_apply_fn() (jaxili.train.TrainerModule method) init_checkpointer() (jaxili.train.TrainerModule method) init_logger() (jaxili.train.TrainerModule method)

init_model() (jaxili.train.TrainerModule method) init_optimizer() (jaxili.train.TrainerModule method) initialize_mask() (jaxili.model.MaskedLinear method) is_new_model_better() (jaxili.train.TrainerModule method)

J

jaxili module jaxili.compressor module jaxili.inference module jaxili.inference.nle module jaxili.inference.npe module jaxili.loss module jaxili.model module

jaxili.posterior module jaxili.posterior.base_posterior module jaxili.posterior.direct_posterior module jaxili.posterior.mcmc_posterior module jaxili.train module jaxili.utils module jaxili.validation module

L

layers (jaxili.model.AffineCoupling attribute) (jaxili.model.ConditionalMAF attribute) (jaxili.model.ConditionalRealNVP attribute) (jaxili.model.MixtureDensityNetwork attribute) load_from_checkpoints() (jaxili.inference.nle.NLE class method) (jaxili.inference.npe.NPE class method) (jaxili.train.TrainerModule class method) load_model() (jaxili.train.TrainerModule method) log_likelihood() (jaxili.posterior.mcmc_posterior.MCMCPosterior method) log_prior() (jaxili.posterior.mcmc_posterior.MCMCPosterior method) log_prob() (jaxili.model.Compressor_w_NDE method) (jaxili.model.ConditionalMAF method) (jaxili.model.ConditionalRealNVP method) (jaxili.model.MixtureDensityNetwork method) (jaxili.model.NDE_Compressor method) (jaxili.model.NDE_w_Standardization method) (jaxili.model.NDENetwork method)

log_prob_compressed() (jaxili.model.NDE_Compressor method) log_prob_from_compressed() (jaxili.model.Compressor_w_NDE method) loss_mmd_npe() (in module jaxili.loss) loss_nll_nle() (in module jaxili.loss) loss_nll_npe() (in module jaxili.loss)

M

MAFLayer (class in jaxili.model) mask (jaxili.model.MaskedLinear attribute) MaskedLinear (class in jaxili.model) maximum_mean_discrepancy() (in module jaxili.loss) MCMCPosterior (class in jaxili.posterior.mcmc_posterior) mean (jaxili.compressor.Standardizer attribute) MixtureDensityNetwork (class in jaxili.model) MLPCompressor (class in jaxili.compressor) mmd_kernel() (in module jaxili.loss) mmd_summary_space() (in module jaxili.loss) module jaxili jaxili.compressor jaxili.inference jaxili.inference.nle jaxili.inference.npe jaxili.loss jaxili.model jaxili.posterior jaxili.posterior.base_posterior jaxili.posterior.direct_posterior jaxili.posterior.mcmc_posterior jaxili.train jaxili.utils jaxili.validation

N

n_components (jaxili.model.MixtureDensityNetwork attribute) n_cond (jaxili.model.ConditionalMADE attribute) (jaxili.model.ConditionalMAF attribute) (jaxili.model.ConditionalRealNVP attribute) (jaxili.model.MAFLayer attribute) (jaxili.model.MixtureDensityNetwork attribute) n_in (jaxili.model.ConditionalMADE attribute) (jaxili.model.ConditionalMAF attribute) (jaxili.model.ConditionalRealNVP attribute) (jaxili.model.MAFLayer attribute) (jaxili.model.MixtureDensityNetwork attribute) n_layers (jaxili.model.ConditionalMAF attribute) (jaxili.model.ConditionalRealNVP attribute) n_out (jaxili.model.MaskedLinear attribute) name (jaxili.compressor.CNN2DCompressor attribute) (jaxili.compressor.Identity attribute) (jaxili.compressor.MLPCompressor attribute) (jaxili.compressor.Standardizer attribute) (jaxili.model.AffineCoupling attribute) (jaxili.model.Compressor_w_NDE attribute) (jaxili.model.ConditionalMADE attribute) (jaxili.model.ConditionalMAF attribute) (jaxili.model.ConditionalRealNVP attribute) (jaxili.model.MAFLayer attribute) (jaxili.model.MaskedLinear attribute) (jaxili.model.MixtureDensityNetwork attribute) (jaxili.model.NDE_Compressor attribute) (jaxili.model.NDE_w_Standardization attribute) (jaxili.model.NDENetwork attribute)

nde (jaxili.model.NDE_Compressor attribute) (jaxili.model.NDE_w_Standardization attribute) NDE_Compressor (class in jaxili.model) nde_hparams (jaxili.model.NDE_Compressor attribute) NDE_w_Standardization (class in jaxili.model) NDENetwork (class in jaxili.model) NeuralPosterior (class in jaxili.posterior.base_posterior) NLE (class in jaxili.inference.nle) NPE (class in jaxili.inference.npe)

O

on_training_epoch_end() (jaxili.train.TrainerModule method) on_training_start() (jaxili.train.TrainerModule method)

on_validation_epoch_end() (jaxili.train.TrainerModule method) output_size (jaxili.compressor.CNN2DCompressor attribute) (jaxili.compressor.MLPCompressor attribute)

P

parent (jaxili.compressor.CNN2DCompressor attribute) (jaxili.compressor.Identity attribute) (jaxili.compressor.MLPCompressor attribute) (jaxili.compressor.Standardizer attribute) (jaxili.model.AffineCoupling attribute) (jaxili.model.Compressor_w_NDE attribute) (jaxili.model.ConditionalMADE attribute) (jaxili.model.ConditionalMAF attribute) (jaxili.model.ConditionalRealNVP attribute) (jaxili.model.MAFLayer attribute) (jaxili.model.MaskedLinear attribute) (jaxili.model.MixtureDensityNetwork attribute) (jaxili.model.NDE_Compressor attribute) (jaxili.model.NDE_w_Standardization attribute) (jaxili.model.NDENetwork attribute)

print_implemented_methods() (jaxili.posterior.mcmc_posterior.MCMCPosterior method) print_tabulate() (jaxili.train.TrainerModule method)

R

random_order (jaxili.model.ConditionalMADE attribute) replace() (jaxili.train.TrainState method)

reverse (jaxili.model.MAFLayer attribute) rng (jaxili.train.TrainState attribute) run_model_init() (jaxili.train.TrainerModule method)

S

sample() (jaxili.model.Compressor_w_NDE method) (jaxili.model.ConditionalMAF method) (jaxili.model.ConditionalRealNVP method) (jaxili.model.MixtureDensityNetwork method) (jaxili.model.NDE_Compressor method) (jaxili.model.NDE_w_Standardization method) (jaxili.model.NDENetwork method) (jaxili.posterior.base_posterior.NeuralPosterior method) (jaxili.posterior.direct_posterior.DirectPosterior method) (jaxili.posterior.mcmc_posterior.MCMCPosterior method) save_metrics() (jaxili.train.TrainerModule method) save_model() (jaxili.train.TrainerModule method) scope (jaxili.compressor.CNN2DCompressor attribute) (jaxili.compressor.Identity attribute) (jaxili.compressor.MLPCompressor attribute) (jaxili.compressor.Standardizer attribute) (jaxili.model.AffineCoupling attribute) (jaxili.model.Compressor_w_NDE attribute) (jaxili.model.ConditionalMADE attribute) (jaxili.model.ConditionalMAF attribute) (jaxili.model.ConditionalRealNVP attribute) (jaxili.model.MAFLayer attribute) (jaxili.model.MaskedLinear attribute) (jaxili.model.MixtureDensityNetwork attribute) (jaxili.model.NDE_Compressor attribute) (jaxili.model.NDE_w_Standardization attribute) (jaxili.model.NDENetwork attribute)

seed (jaxili.model.ConditionalMADE attribute) (jaxili.model.ConditionalMAF attribute) (jaxili.model.MAFLayer attribute) set_dataloader() (jaxili.inference.nle.NLE method) (jaxili.inference.npe.NPE method) set_dataset() (jaxili.inference.nle.NLE method) (jaxili.inference.npe.NPE method) set_default_x() (jaxili.posterior.base_posterior.NeuralPosterior method) (jaxili.posterior.mcmc_posterior.MCMCPosterior method) set_loss_fn() (jaxili.inference.nle.NLE method) (jaxili.inference.npe.NPE method) set_mcmc_kwargs() (jaxili.posterior.mcmc_posterior.MCMCPosterior method) set_mcmc_method() (jaxili.posterior.mcmc_posterior.MCMCPosterior method) set_model_hparams() (jaxili.inference.nle.NLE method) (jaxili.inference.npe.NPE method) set_prior() (jaxili.posterior.mcmc_posterior.MCMCPosterior method) setup() (jaxili.model.ConditionalMADE method) (jaxili.model.ConditionalMAF method) (jaxili.model.NDE_Compressor method) standardize() (jaxili.model.NDE_w_Standardization method) Standardizer (class in jaxili.compressor) std (jaxili.compressor.Standardizer attribute)

T

tracker() (jaxili.train.TrainerModule method) train() (jaxili.inference.nle.NLE method) (jaxili.inference.npe.NPE method) train_epoch() (jaxili.train.TrainerModule method)

train_model() (jaxili.train.TrainerModule method) TrainerModule (class in jaxili.train) TrainState (class in jaxili.train) transformation (jaxili.model.NDE_w_Standardization attribute)

U

unnormalized_log_prob() (jaxili.posterior.base_posterior.NeuralPosterior method) (jaxili.posterior.direct_posterior.DirectPosterior method) (jaxili.posterior.mcmc_posterior.MCMCPosterior method)

unstandardize() (jaxili.model.NDE_w_Standardization method) use_reverse (jaxili.model.ConditionalMAF attribute)

V

validate_theta_x() (in module jaxili.utils)

W

write_config() (jaxili.train.TrainerModule method)

Y

y (jaxili.model.AffineCoupling attribute)