jaxili package

JaxILI: A JAX library for Implicit Likelihood Inference using Neural Density Estimation.

Subpackages

• jaxili.inference package
  • Submodules
    • jaxili.inference.nle module
      • NLE
    • jaxili.inference.npe module
      • NPE
• jaxili.posterior package
  • Submodules
    • jaxili.posterior.base_posterior module
      • NeuralPosterior
    • jaxili.posterior.direct_posterior module
      • DirectPosterior
    • jaxili.posterior.mcmc_posterior module
      • MCMCPosterior

Submodules

• jaxili.compressor module
  • CNN2DCompressor
    • CNN2DCompressor.activation
    • CNN2DCompressor.name
    • CNN2DCompressor.output_size
    • CNN2DCompressor.parent
    • CNN2DCompressor.scope
  • Identity
    • Identity.name
    • Identity.parent
    • Identity.scope
  • MLPCompressor
    • MLPCompressor.activation
    • MLPCompressor.hidden_size
    • MLPCompressor.name
    • MLPCompressor.output_size
    • MLPCompressor.parent
    • MLPCompressor.scope
  • Standardizer
    • Standardizer.mean
    • Standardizer.name
    • Standardizer.parent
    • Standardizer.scope
    • Standardizer.std
• jaxili.loss module
  • gaussian_kernel_matrix()
  • loss_mmd_npe()
  • loss_nll_nle()
  • loss_nll_npe()
  • maximum_mean_discrepancy()
  • mmd_kernel()
  • mmd_summary_space()
• jaxili.model module
  • AffineCoupling
    • AffineCoupling.activation
    • AffineCoupling.layers
    • AffineCoupling.name
    • AffineCoupling.parent
    • AffineCoupling.scope
    • AffineCoupling.y
  • Compressor_w_NDE
    • Compressor_w_NDE.compress()
    • Compressor_w_NDE.log_prob()
    • Compressor_w_NDE.log_prob_from_compressed()
    • Compressor_w_NDE.name
    • Compressor_w_NDE.parent
    • Compressor_w_NDE.sample()
    • Compressor_w_NDE.scope
  • ConditionalMADE
    • ConditionalMADE.activation
    • ConditionalMADE.gaussian
    • ConditionalMADE.hidden_dims
    • ConditionalMADE.n_cond
    • ConditionalMADE.n_in
    • ConditionalMADE.name
    • ConditionalMADE.parent
    • ConditionalMADE.random_order
    • ConditionalMADE.scope
    • ConditionalMADE.seed
    • ConditionalMADE.setup()
  • ConditionalMAF
    • ConditionalMAF.activation
    • ConditionalMAF.backward()
    • ConditionalMAF.layers
    • ConditionalMAF.log_prob()
    • ConditionalMAF.n_cond
    • ConditionalMAF.n_in
    • ConditionalMAF.n_layers
    • ConditionalMAF.name
    • ConditionalMAF.parent
    • ConditionalMAF.sample()
    • ConditionalMAF.scope
    • ConditionalMAF.seed
    • ConditionalMAF.setup()
    • ConditionalMAF.use_reverse
  • ConditionalRealNVP
    • ConditionalRealNVP.activation
    • ConditionalRealNVP.layers
    • ConditionalRealNVP.log_prob()
    • ConditionalRealNVP.n_cond
    • ConditionalRealNVP.n_in
    • ConditionalRealNVP.n_layers
    • ConditionalRealNVP.name
    • ConditionalRealNVP.parent
    • ConditionalRealNVP.sample()
    • ConditionalRealNVP.scope
  • MAFLayer
    • MAFLayer.activation
    • MAFLayer.backward()
    • MAFLayer.forward()
    • MAFLayer.hidden_dims
    • MAFLayer.n_cond
    • MAFLayer.n_in
    • MAFLayer.name
    • MAFLayer.parent
    • MAFLayer.reverse
    • MAFLayer.scope
    • MAFLayer.seed
  • MaskedLinear
    • MaskedLinear.bias
    • MaskedLinear.initialize_mask()
    • MaskedLinear.mask
    • MaskedLinear.n_out
    • MaskedLinear.name
    • MaskedLinear.parent
    • MaskedLinear.scope
  • MixtureDensityNetwork
    • MixtureDensityNetwork.activation
    • MixtureDensityNetwork.layers
    • MixtureDensityNetwork.log_prob()
    • MixtureDensityNetwork.n_components
    • MixtureDensityNetwork.n_cond
    • MixtureDensityNetwork.n_in
    • MixtureDensityNetwork.name
    • MixtureDensityNetwork.parent
    • MixtureDensityNetwork.sample()
    • MixtureDensityNetwork.scope
  • NDENetwork
    • NDENetwork.log_prob()
    • NDENetwork.name
    • NDENetwork.parent
    • NDENetwork.sample()
    • NDENetwork.scope
  • NDE_Compressor
    • NDE_Compressor.compressor
    • NDE_Compressor.compressor_hparams
    • NDE_Compressor.log_prob()
    • NDE_Compressor.log_prob_compressed()
    • NDE_Compressor.name
    • NDE_Compressor.nde
    • NDE_Compressor.nde_hparams
    • NDE_Compressor.parent
    • NDE_Compressor.sample()
    • NDE_Compressor.scope
    • NDE_Compressor.setup()
  • NDE_w_Standardization
    • NDE_w_Standardization.embedding()
    • NDE_w_Standardization.embedding_net
    • NDE_w_Standardization.log_prob()
    • NDE_w_Standardization.name
    • NDE_w_Standardization.nde
    • NDE_w_Standardization.parent
    • NDE_w_Standardization.sample()
    • NDE_w_Standardization.scope
    • NDE_w_Standardization.standardize()
    • NDE_w_Standardization.transformation
    • NDE_w_Standardization.unstandardize()
• jaxili.train module
  • TrainState
    • TrainState.batch_stats
    • TrainState.replace()
    • TrainState.rng
  • TrainerModule
    • TrainerModule.bind_model()
    • TrainerModule.create_functions()
    • TrainerModule.create_jitted_functions()
    • TrainerModule.eval_model()
    • TrainerModule.generate_config()
    • TrainerModule.init_apply_fn()
    • TrainerModule.init_checkpointer()
    • TrainerModule.init_logger()
    • TrainerModule.init_model()
    • TrainerModule.init_optimizer()
    • TrainerModule.is_new_model_better()
    • TrainerModule.load_from_checkpoints()
    • TrainerModule.load_model()
    • TrainerModule.on_training_epoch_end()
    • TrainerModule.on_training_start()
    • TrainerModule.on_validation_epoch_end()
    • TrainerModule.print_tabulate()
    • TrainerModule.run_model_init()
    • TrainerModule.save_metrics()
    • TrainerModule.save_model()
    • TrainerModule.tracker()
    • TrainerModule.train_epoch()
    • TrainerModule.train_model()
    • TrainerModule.write_config()
• jaxili.utils module
  • check_density_estimator()
  • check_hparams_maf()
  • check_hparams_mdn()
  • check_hparams_realnvp()
  • create_data_loader()
  • handle_non_serializable()
  • validate_theta_x()
• jaxili.validation module