API Reference

This page provides an auto-generated summary of NeuPI’s public API, generated directly from the source code’s docstrings.

neupi

NeuPI: Neural-Probabilistic Inference Library

neupi.core

neupi.embedding

NeuPI: Neural-Probabilistic Inference Library

class neupi.DiscreteEmbedder(num_vars: int, query_embedding: float = 0.0, unobserved_embedding: float = 1.0)[source]

Bases: Embedding

Creates feature embeddings from variable assignments and bucket information.

This preprocessor converts a tensor of binary variable assignments into a new feature space suitable for input to a neural network. It expands each variable v into a two-dimensional representation [v, 1-v] and then replaces the features for query and unobserved variables with specific embedding values.

Parameters:

  • query_embedding (float) – The value to use for query variables. Defaults to 0.0.

  • unobserved_embedding (float) – The value to use for unobserved variables. Defaults to 1.0.

class neupi.ITSELF_Engine(model: torch.nn.Module, pgm_evaluator: torch.nn.Module, loss_fn: callable, optimizer_cls: Type[torch.optim.Optimizer], discretizer: torch.nn.Module, refinement_lr: float, refinement_steps: int, device: str = 'cpu')[source]

Bases: BaseInferenceModule

Handles Inference Time Self-Supervised Training (ITSELF).

For each batch of data, this engine performs a few steps of fine-tuning on a copy of the model to refine the predictions for that specific batch.

Parameters:

  • model (torch.nn.Module) – The base, pre-trained neural network model.

  • pgm_evaluator (torch.nn.Module) – The PGM evaluator for loss calculation.

  • loss_fn (callable) – The loss function used during refinement.

  • optimizer_cls (Type[torch.optim.Optimizer]) – The class of the optimizer to use for refinement (e.g., torch.optim.Adam).

  • refinement_lr (float) – The learning rate for the test-time optimizer.

  • refinement_steps (int) – The number of optimization steps to perform per instance.

  • device (str) – The device to run inference on (‘cpu’ or ‘cuda’).

References

Arya, S., Rahman, T., & Gogate, V. G. (2024). A neural network approach for efficiently answering most probable explanation queries in probabilistic models. NeurIPS 2024. https://openreview.net/forum?id=ufPPf9ghzP

run(data) → Dict[str, torch.Tensor][source]

Performs test-time refinement for each batch in the dataloader or a single batch. :param data: Either a DataLoader or a tuple of (evidence_data, evidence_mask, query_mask, unobs_mask).

Returns:

Final assignments.

Return type:

Dict[str, torch.Tensor]

class neupi.IdentityEmbedding(num_vars: int)[source]

Bases: Embedding

An identity embedding that does not change the input.

class neupi.KNearestDiscretizer(*args: Any, **kwargs: Any)[source]

Bases: BaseDiscretizer

Finds k-nearest binary vectors for query variables using a scoring function.

This method generates candidate binary assignments for query variables that are “close” to the continuous predictions, scores them using the PGM evaluator, and selects the best one.

Parameters:

  • pgm_evaluator (Callable) – The PGM evaluator (e.g., MarkovNetwork) which acts as the scoring function.

  • k (int) – Number of nearest binary vectors to consider.

  • batch_size (int) – Batch size for scoring candidate assignments.

References

Arya, Shivvrat, Rahman, Tahrima, and Gogate, Vibhav Giridhar. SINE: Scalable MPE Inference for Probabilistic Graphical Models Using Advanced Neural Embeddings. Proceedings of the 28th International Conference on Artificial Intelligence and Statistics (AISTATS), 2025.

class neupi.MADE(*args: Any, **kwargs: Any)[source]

Bases: Module

evaluate(x)[source]
forward(x)[source]
update_masks()[source]
class neupi.MLP(*args: Any, **kwargs: Any)[source]

Bases: BaseNNModel

A flexible Multi-Layer Perceptron (MLP) network.

This module creates a fully connected neural network with configurable hidden layers, activation functions, batch normalization, and dropout. We extract the input size from the embedding.

Parameters:

  • hidden_sizes (List[int]) – A list where each element is the number of neurons in a hidden layer.

  • output_size (int) – The number of neurons in the output layer.

  • hidden_activation (str) – The activation function for hidden layers. Supported: ‘relu’, ‘leaky_relu’.

  • use_batchnorm (bool) – If True, adds a BatchNorm1d layer after each hidden activation. Defaults to True.

  • dropout_rate (float) – The dropout probability. If 0.0, no dropout is applied. Defaults to 0.0.

forward(evidence_data: torch.Tensor, evidence_mask: torch.Tensor, query_mask: torch.Tensor, unobs_mask: torch.Tensor) → torch.Tensor[source]

Defines the forward pass of the MLP.

initialize_weights(nonlinearity: str = 'relu')[source]

Initializes the weights of the network using appropriate methods.

class neupi.MarkovNetwork(*args: Any, **kwargs: Any)[source]

Bases: BaseProbModel

A PyTorch module to evaluate the log-likelihood of assignments in a binary Markov Network.

This class loads a Markov Network from a .uai file and provides an efficient, vectorized method to compute the log-likelihood for a batch of variable assignments. It supports both pairwise and higher-order factor models.

Parameters:

  • uai_file (str) – Path to the .uai file defining the Markov Network.

  • device (str or torch.device) – The device to perform computations on (‘cpu’ or ‘cuda’).

Note: For now, models from the UAI format are supported. Other formats can be supported in by rewriting the UAIParser class.

Example: ` from neupi.pm.pgm.mn import MarkovNetwork mn = MarkovNetwork("path/to/mn.uai") `

forward(x: torch.Tensor) → torch.Tensor[source]

Alias for the evaluate method.

Parameters:

x (torch.Tensor) – A binary tensor of assignments of shape (batch_size, num_variables).

Returns:

A tensor of log-likelihood scores of shape (batch_size,).

Return type:

torch.Tensor

class neupi.OAUAI(*args: Any, **kwargs: Any)[source]

Bases: BaseDiscretizer

Discretizes by performing an exhaustive search over the k most uncertain variables.

This method identifies the ‘k’ query variables with probabilities closest to 0.5, generates all 2^k possible binary assignments for this subset, scores each one using the PGM evaluator, and selects the best assignment.

Note: This is a naive oracle that looks at all 2^k possible assignments and selects the best one. Other oracles can also be used (such as daoopt).

Parameters:

  • pgm_evaluator (Callable) – The PGM evaluator (e.g., MarkovNetwork) which acts as the scoring function.

  • k (int) – The number of most uncertain query variables to search over.

  • threshold (float) – The baseline threshold for comparison. Defaults to 0.5.

References

Arya, Shivvrat, Rahman, Tahrima, and Gogate, Vibhav Giridhar. SINE: Scalable MPE Inference for Probabilistic Graphical Models Using Advanced Neural Embeddings. Proceedings of the 28th International Conference on Artificial Intelligence and Statistics (AISTATS), 2025.

class neupi.SelfSupervisedTrainer(model: torch.nn.Module, pgm_evaluator: torch.nn.Module, loss_fn: callable, optimizer: torch.optim.Optimizer, device: str = 'cpu')[source]

Bases: BaseTrainer

A trainer for self-supervised learning of neural PGM solvers.

This class handles the training and validation loops, including the forward pass, loss calculation, backpropagation, and optimizer steps.

Parameters:

  • model (torch.nn.Module) – The neural network model to be trained.

  • pgm_evaluator (torch.nn.Module) – The PGM evaluator (e.g., MarkovNetwork) used for calculating log-likelihoods.

  • loss_fn (callable) – The loss function. It should accept predictions and the PGM evaluator as arguments.

  • optimizer (torch.optim.Optimizer) – The optimizer for updating model weights.

  • device (str) – The device to run training on (‘cpu’ or ‘cuda’).

fit(dataloader: torch.utils.data.DataLoader, num_epochs: int)[source]

Runs the full training loop for a specified number of epochs.

Parameters:

  • dataloader (DataLoader) – The DataLoader providing training data.

  • num_epochs (int) – The number of epochs to train for.

step(batch_data)[source]

Performs a single training step on a batch of data.

class neupi.SinglePassInferenceEngine(model: torch.nn.Module, discretizer: torch.nn.Module, device: str = 'cpu')[source]

Bases: BaseInferenceModule

Handles the inference process for a trained neural PGM solver.

This class takes a trained model and runs it on a dataset, collecting the outputs. It operates in no_grad mode for efficiency.

Parameters:

model (torch.nn.Module) – The trained neural network model.

References

Arya, S., Rahman, T., & Gogate, V. (2025). SINE: Scalable MPE inference for probabilistic graphical models using advanced neural embeddings. AISTATS 2025. https://openreview.net/forum?id=6tZxEVlpnL

Arya, S., Rahman, T., & Gogate, V. (2024). Learning to Solve the Constrained Most Probable Explanation Task in Probabilistic Graphical Models. Proceedings of The 27th International Conference on Artificial Intelligence and Statistics. International Conference on Artificial Intelligence and Statistics, PMLR, pp. 2791–2799. https://proceedings.mlr.press/v238/arya24b.html

Arya, S., Rahman, T., & Gogate, V. (2024). Neural Network Approximators for Marginal MAP in Probabilistic Circuits. Proceedings of the AAAI Conference on Artificial Intelligence, 38(10), 10918–10926. https://doi.org/10.1609/aaai.v38i10.28966

run(data) → Dict[str, torch.Tensor]

Performs a forward pass on the entire dataset or a single batch to get predictions.

Parameters:

data – Either a DataLoader or a tuple/list of (evidence_data, evidence_mask, query_mask, unobs_mask).

Returns:

Concatenated tensors for:

  • ’raw_outputs’: Raw logits from the network.

  • ’prob_outputs’: Probabilities after sigmoid activation.

  • ’final_assignments’: Final assignments after applying evidence.

Return type:

Dict[str, torch.Tensor]

class neupi.SumProductNetwork(*args: Any, **kwargs: Any)[source]

Bases: BaseProbModel

evaluate(x)[source]

Evaluate the SPN model on input data.

Parameters:

x – Tensor representing input data (batch_size, input_size).

Returns:

Function value at the root of the SPN.

forward(x)[source]
class neupi.ThresholdDiscretizer(*args: Any, **kwargs: Any)[source]

Bases: BaseDiscretizer

Discretizes a tensor of probabilities to binary assignments based on a threshold.

This class acts as a callable function. An instance can be passed to an inference engine and called to perform the discretization.

Parameters:

threshold (float) – The threshold value. Values >= threshold will be 1, and values < threshold will be 0. Defaults to 0.5.

References

Arya, S., Rahman, T., & Gogate, V. (2024). Learning to Solve the Constrained Most Probable Explanation Task in Probabilistic Graphical Models. Proceedings of The 27th International Conference on Artificial Intelligence and Statistics. International Conference on Artificial Intelligence and Statistics, PMLR, pp. 2791–2799. https://proceedings.mlr.press/v238/arya24b.html

Arya, S., Rahman, T., & Gogate, V. (2024). Neural Network Approximators for Marginal MAP in Probabilistic Circuits. Proceedings of the AAAI Conference on Artificial Intelligence, 38(10), 10918–10926. https://doi.org/10.1609/aaai.v38i10.28966

Arya, S., Rahman, T., & Gogate, V. G. (2024). A neural network approach for efficiently answering most probable explanation queries in probabilistic models. NeurIPS 2024. https://openreview.net/forum?id=ufPPf9ghzP

neupi.factory(kind: str, name: str, *args: Any, **kwargs: Any) → Any

Instantiate a class from the registry by its kind and name.

neupi.mpe_log_likelihood_loss(predictions: torch.Tensor, pgm_evaluator: torch.nn.Module, aggregate: str = 'avg') → torch.Tensor[source]

Calculates the negative log-likelihood for a batch of MPE predictions.

This loss function is designed for training neural networks to solve MPE inference. The goal is to maximize the log-likelihood of the predicted assignments. Since optimizers minimize loss, this function returns the negative of the summed log-likelihoods.

Note: The same loss functions can be used for MMAP inference over PCs (no other PMs are supported yet for MMAP). For MMAP, make the unobserved variables -1s and the correct scores will be computed by the evaluator.

Parameters:

  • predictions (torch.Tensor) – A batch of binary assignments (0s or 1s) predicted by a neural network. Shape: (batch_size, num_variables).

  • pgm_evaluator (torch.nn.Module) – An instantiated PGM evaluator object from neupi.pgm (e.g., MarkovNetwork, SumProductNetwork) that has a forward or evaluate method.

Returns:

A scalar tensor representing the total negative log-likelihood,

to be used for backpropagation.

Return type:

torch.Tensor