Leaf Modules¶

Probability distributions at the leaves of the probabilistic circuit. See Base Classes for the abstract LeafModule base class.

Continuous Distributions¶

Normal¶

Gaussian distribution with learnable mean and standard deviation.

class spflow.modules.leaves.Normal(scope, out_channels=None, num_repetitions=1, parameter_fn=None, validate_args=True, loc=None, scale=None)[source]¶

Bases: LeafModule

Normal (Gaussian) distribution leaf module.

Parameterized by mean μ and standard deviation σ (stored in log-space).

loc¶: Mean parameter.

std¶: Standard deviation (accessed via property, stored as log_std).

__init__(scope, out_channels=None, num_repetitions=1, parameter_fn=None, validate_args=True, loc=None, scale=None)[source]¶

Initialize Normal distribution.

Parameters:

scope – Variable scope (Scope, int, or list[int]).
out_channels (int) – Number of output channels (inferred from params if None).
num_repetitions (int) – Number of repetitions (for 3D event shapes).
parameter_fn (Module) – Optional neural network for parameter generation.
loc (Tensor) – Mean tensor μ.
scale (Tensor) – Standard deviation tensor σ > 0.

params()[source]¶: Returns the parameters of the distribution.

property scale: Tensor¶: Standard deviation in natural space (read via exp of log_std).

LogNormal¶

Log-normal distribution (exponential of a normal distribution).

class spflow.modules.leaves.LogNormal(scope, out_channels=None, num_repetitions=1, parameter_fn=None, validate_args=True, loc=None, scale=None)[source]¶

Bases: LeafModule

Log-Normal distribution leaf for modeling positive-valued data.

Note: Parameters μ and σ apply to ln(x), not x itself. Standard deviation σ is stored in log-space for numerical stability.

mean¶: Mean μ of log-space distribution.

std¶: Standard deviation σ > 0 of log-space distribution (via exp of log_std).

distribution¶: Underlying torch.distributions.LogNormal.

__init__(scope, out_channels=None, num_repetitions=1, parameter_fn=None, validate_args=True, loc=None, scale=None)[source]¶

Initialize Log-Normal distribution.

Parameters:

scope – Variable scope (Scope, int, or list[int]).
out_channels (int) – Number of output channels (inferred from params if None).
num_repetitions (int) – Number of repetitions (for 3D event shapes).
parameter_fn (Module) – Optional neural network for parameter generation.
validate_args (bool | None) – Whether to enable torch.distributions argument validation.
loc (Tensor) – Mean μ of log-space distribution.
scale (Tensor) – Standard deviation σ > 0 of log-space distribution.

params()[source]¶

Returns distribution parameters.

Return type:: dict[str, Tensor]

property scale: Tensor¶: Standard deviation in natural space (read via exp of log_scale).

Exponential¶

Exponential distribution with learnable rate parameter.

class spflow.modules.leaves.Exponential(scope, out_channels=None, num_repetitions=1, parameter_fn=None, validate_args=True, rate=None)[source]¶

Bases: LeafModule

Exponential distribution leaf for modeling time-between-events.

Parameterized by rate λ > 0 (stored in log-space for numerical stability).

rate¶: Rate parameter λ (accessed via property, stored as log_rate).

distribution¶: Underlying torch.distributions.Exponential.

__init__(scope, out_channels=None, num_repetitions=1, parameter_fn=None, validate_args=True, rate=None)[source]¶

Initialize Exponential distribution leaf.

Parameters:

scope – Variable scope (Scope, int, or list[int]).
out_channels (int) – Number of output channels (inferred from params if None).
num_repetitions (int) – Number of repetitions (for 3D event shapes).
parameter_fn (Module) – Optional neural network for parameter generation.
validate_args (bool | None) – Whether to enable torch.distributions argument validation.
rate (Tensor) – Rate parameter λ > 0.

params()[source]¶

Returns distribution parameters.

Return type:: dict[str, Tensor]

property rate: Tensor¶: Rate parameter in natural space (read via exp of log_rate).

Gamma¶

Gamma distribution with learnable alpha and beta parameters.

class spflow.modules.leaves.Gamma(scope, out_channels=None, num_repetitions=1, parameter_fn=None, validate_args=True, concentration=None, rate=None)[source]¶

Bases: LeafModule

Gamma distribution leaf for modeling positive-valued continuous data.

Parameterized by shape α > 0 and rate β > 0 (both stored in log-space for numerical stability).

alpha¶: Shape parameter α (accessed via property, stored as log_alpha).

beta¶: Rate parameter β (accessed via property, stored as log_beta).

distribution¶: Underlying torch.distributions.Gamma.

__init__(scope, out_channels=None, num_repetitions=1, parameter_fn=None, validate_args=True, concentration=None, rate=None)[source]¶

Initialize Gamma distribution leaf.

Parameters:

scope – Variable scope (Scope, int, or list[int]).
out_channels (int) – Number of output channels (inferred from params if None).
num_repetitions (int) – Number of repetitions (for 3D event shapes).
parameter_fn (Module) – Optional neural network for parameter generation.
validate_args (bool | None) – Whether to enable torch.distributions argument validation.
concentration (Tensor) – Shape parameter α > 0.
rate (Tensor) – Rate parameter β > 0.

conditional_distribution(evidence)[source]¶

Generates torch.distributions object conditionally based on evidence.

Parameters:: evidence (Tensor) – Evidence tensor for conditioning.
Return type:: Gamma
Returns:: torch.distributions.Distribution constructed from conditional parameters.

params()[source]¶

Returns distribution parameters.

Return type:: dict[str, Tensor]

property concentration: Tensor¶: Shape parameter in natural space (read via exp of log_alpha).

property rate: Tensor¶: Rate parameter in natural space (read via exp of log_beta).

Uniform¶

Uniform distribution over a bounded interval.

class spflow.modules.leaves.Uniform(scope, out_channels=None, num_repetitions=1, low=None, high=None, validate_args=True, support_outside=True)[source]¶

Bases: LeafModule

Uniform distribution leaf with fixed interval bounds.

Note: Interval bounds are fixed buffers and cannot be learned.

start¶: Start of interval (fixed buffer).

end¶: End of interval (fixed buffer).

end_next¶: Next representable value after end.

support_outside¶: Whether values outside [start, end] are supported.

distribution¶: Underlying torch.distributions.Uniform.

__init__(scope, out_channels=None, num_repetitions=1, low=None, high=None, validate_args=True, support_outside=True)[source]¶

Initialize Uniform distribution leaf.

Parameters:

scope (Scope) – Variable scope.
out_channels (int | None) – Number of output channels (inferred from params if None).
num_repetitions (int) – Number of repetitions.
low (Tensor | None) – Lower bound tensor (required).
high (Tensor | None) – Upper bound tensor (required).
validate_args (bool | None) – Whether to enable torch.distributions argument validation.
support_outside (bool) – Whether values outside [start, end] are supported.

expectation_maximization(data, bias_correction=False, cache=None)[source]¶

Perform single EM step.

Parameters:

data (Tensor) – Input data tensor.
cache (Cache | None) – Optional cache dictionary.

Return type:

None

params()[source]¶

Returns distribution parameters.

Return type:: dict[str, Tensor]

property mode: Tensor¶: Returns the mode (midpoint) of the distribution.

Discrete Distributions¶

Categorical¶

Categorical distribution over K categories with learnable probabilities.

class spflow.modules.leaves.Categorical(scope, out_channels=None, num_repetitions=1, K=None, probs=None, logits=None, parameter_fn=None, validate_args=True)[source]¶

Bases: LeafModule

Categorical distribution leaf for discrete choice over K categories.

p¶: Categorical probabilities (normalized, includes extra dimension for K).

K¶: Number of categories.

distribution¶: Underlying torch.distributions.Categorical.

__init__(scope, out_channels=None, num_repetitions=1, K=None, probs=None, logits=None, parameter_fn=None, validate_args=True)[source]¶

Initialize Categorical distribution leaf module.

Parameters:

scope (Scope) – The scope of the distribution.
out_channels (int | None) – Number of output channels (inferred from params if None).
num_repetitions (int) – Number of repetitions for the distribution.
K (int | None) – Number of categories (optional if parameter tensor provided).
probs (Tensor | None) – Probability tensor over categories.
logits (Tensor | None) – Logits tensor over categories.
parameter_fn (Module | None) – Optional neural network for parameter generation.
validate_args (bool | None) – Whether to enable torch.distributions argument validation.

params()[source]¶

Returns distribution parameters.

Return type:: dict[str, Tensor]

property logits: Tensor¶: Logits directly parameterizing the categorical distribution.

property probs: Tensor¶: Categorical probabilities in natural space (read via softmax of logits).

Bernoulli¶

Bernoulli distribution (binary outcomes) with learnable probability.

class spflow.modules.leaves.Bernoulli(scope, out_channels=None, num_repetitions=1, parameter_fn=None, validate_args=True, probs=None, logits=None)[source]¶

Bases: LeafModule

Bernoulli distribution leaf module.

Binary random variable with success probability p ∈ [0, 1]. Parameterized by success probability p ∈ [0, 1] (stored in logit-space for numerical stability).

p¶: Success probability (BoundedParameter).

distribution¶: Underlying torch.distributions.Bernoulli.

__init__(scope, out_channels=None, num_repetitions=1, parameter_fn=None, validate_args=True, probs=None, logits=None)[source]¶

Initialize Bernoulli distribution.

Parameters:

scope – Variable scope (Scope, int, or list[int]).
out_channels (int) – Number of output channels (inferred from params if None).
num_repetitions (int) – Number of repetitions (for 3D event shapes).
parameter_fn (Module) – Optional neural network for parameter generation.
validate_args (bool | None) – Whether to enable torch.distributions argument validation.
probs (Tensor | None) – Success probability tensor in [0, 1].
logits (Tensor | None) – Logits corresponding to success probability.

params()[source]¶

Returns distribution parameters.

Return type:: dict[str, Tensor]

property logits: Tensor¶: Logits of the Bernoulli distribution.

property probs: Tensor¶: Success probability in natural space (read via inverse projection of logits).

Binomial¶

Binomial distribution with learnable probability and fixed number of trials.

class spflow.modules.leaves.Binomial(scope, out_channels=None, num_repetitions=1, total_count=None, probs=None, logits=None, parameter_fn=None, validate_args=True)[source]¶

Bases: LeafModule

Binomial distribution leaf module for probabilistic circuits.

Implements univariate Binomial distributions as leaf nodes in probabilistic circuits. Supports parameter learning through maximum likelihood estimation and efficient inference through PyTorch’s built-in distributions.

The Binomial distribution models the number of successes in a fixed number of independent Bernoulli trials, with probability mass function:

P(X = k | n, p) = C(n, k) * p^k * (1-p)^(n-k)

where n is the number of trials (fixed), p is the success probability (learnable, stored in logit-space for numerical stability), and k is the number of successes (0 ≤ k ≤ n).

p¶: Success probability parameter(s) in [0, 1] (BoundedParameter).

n¶: Number of trials parameter(s), non-negative integers (fixed buffer).

distribution¶: Underlying torch.distributions.Binomial.

__init__(scope, out_channels=None, num_repetitions=1, total_count=None, probs=None, logits=None, parameter_fn=None, validate_args=True)[source]¶

Initialize Binomial distribution leaf module.

Parameters:

scope (Scope) – Scope object specifying the scope of the distribution.
out_channels (int) – Number of output channels (inferred from params if None).
num_repetitions (int) – Number of repetitions for the distribution.
total_count (Tensor | None) – Number of trials tensor (required).
probs (Tensor | None) – Success probability tensor (optional, randomly initialized if None).
logits (Tensor | None) – Log-odds tensor for success probability.
parameter_fn (Module) – Optional neural network for parameter generation.
validate_args (bool | None) – Whether to enable torch.distributions argument validation.

params()[source]¶

Returns distribution parameters.

Return type:: dict[str, Tensor]

property logits: Tensor¶: Logits for the success probability.

property probs: Tensor¶: Success probability in natural space (read via inverse projection of logit_p).

property total_count: Tensor¶: Returns the number of trials.

Poisson¶

Poisson distribution with learnable rate parameter.

class spflow.modules.leaves.Poisson(scope, out_channels=None, num_repetitions=1, parameter_fn=None, validate_args=True, rate=None)[source]¶

Bases: LeafModule

Poisson distribution leaf for modeling event counts.

Parameterized by rate λ > 0 (stored in log-space for numerical stability).

rate¶: Rate parameter λ (stored as log_rate internally).

distribution¶: Underlying torch.distributions.Poisson.

__init__(scope, out_channels=None, num_repetitions=1, parameter_fn=None, validate_args=True, rate=None)[source]¶

Initialize Poisson leaf.

Parameters:

scope – Variable scope (Scope, int, or list[int]).
out_channels (int) – Number of output channels (inferred from params if None).
num_repetitions (int) – Number of repetitions (for 3D event shapes).
parameter_fn (Module) – Optional neural network for parameter generation.
validate_args (bool | None) – Whether to enable torch.distributions argument validation.
rate (Tensor) – Rate parameter λ > 0.

params()[source]¶

Returns distribution parameters.

Return type:: dict[str, Tensor]

property rate: Tensor¶: Rate parameter in natural space (read via exp of log_rate).

Geometric¶

Geometric distribution with learnable success probability.

class spflow.modules.leaves.Geometric(scope, out_channels=None, num_repetitions=1, parameter_fn=None, validate_args=True, probs=None, logits=None)[source]¶

Bases: LeafModule

Geometric distribution leaf for modeling trials until first success.

Parameterized by success probability p ∈ (0, 1] (stored in logit-space for numerical stability).

p¶: Success probability (BoundedParameter).

distribution¶: Underlying torch.distributions.Geometric.

__init__(scope, out_channels=None, num_repetitions=1, parameter_fn=None, validate_args=True, probs=None, logits=None)[source]¶

Initialize Geometric distribution.

Parameters:

scope – Variable scope (Scope, int, or list[int]).
out_channels – Number of output channels (inferred from params if None).
num_repetitions – Number of repetitions (for 3D event shapes).
parameter_fn – Optional neural network for parameter generation.
validate_args (bool | None) – Whether to enable torch.distributions argument validation.
probs (Tensor | None) – Success probability tensor.
logits (Tensor | None) – Log-odds tensor of the success probability.

params()[source]¶

Returns distribution parameters.

Return type:: dict[str, Tensor]

property logits: Tensor¶: Logits for the success probability.

property probs: Tensor¶: Success probability in natural space (read via inverse projection of logits).

NegativeBinomial¶

Negative binomial distribution with learnable parameters.

class spflow.modules.leaves.NegativeBinomial(scope, out_channels=None, num_repetitions=1, total_count=None, probs=None, logits=None, parameter_fn=None, validate_args=True)[source]¶

Bases: LeafModule

Negative Binomial distribution leaf for modeling failures before r-th success.

Note: Parameter n (number of successes) is fixed and cannot be learned. Success probability p is stored in logit-space for numerical stability.

n¶: Fixed number of required successes (buffer).

p¶: Success probability in [0, 1] (stored in logit-space).

distribution¶: Underlying torch.distributions.NegativeBinomial.

__init__(scope, out_channels=None, num_repetitions=1, total_count=None, probs=None, logits=None, parameter_fn=None, validate_args=True)[source]¶

Initialize Negative Binomial distribution leaf module.

Parameters:

scope (Scope) – Scope object specifying the scope of the distribution.
out_channels (int) – Number of output channels (inferred from params if None).
num_repetitions (int) – Number of repetitions for the distribution.
total_count (Tensor | None) – Number of required successes tensor (required).
probs (Tensor | None) – Success probability tensor (optional).
logits (Tensor | None) – Logits of the success probability.
parameter_fn (Module) – Optional neural network for parameter generation.
validate_args (bool | None) – Whether to enable torch.distributions argument validation.

params()[source]¶

Returns distribution parameters.

Return type:: dict[str, Tensor]

property logits: Tensor¶: Logits for success probability.

property probs: Tensor¶: Success probability in natural space (read via inverse projection of logit_p).

property total_count: Tensor¶: Returns the fixed number of required successes.

Hypergeometric¶

Hypergeometric distribution.

class spflow.modules.leaves.Hypergeometric(scope, out_channels=None, num_repetitions=1, K=None, N=None, n=None, validate_args=True)[source]¶

Bases: LeafModule

Hypergeometric distribution leaf for sampling without replacement.

All parameters (K, N, n) are fixed buffers and cannot be learned.

K¶: Number of success states in population (fixed buffer).

N¶: Population size (fixed buffer).

n¶: Number of draws (fixed buffer).

distribution¶: Underlying custom Hypergeometric distribution.

__init__(scope, out_channels=None, num_repetitions=1, K=None, N=None, n=None, validate_args=True)[source]¶

Initialize Hypergeometric distribution leaf module.

Parameters:

scope (Scope) – Scope object specifying the scope of the distribution.
out_channels (int) – Number of output channels (inferred from params if None).
num_repetitions (int) – Number of repetitions for the distribution.
K (Tensor | None) – Number of success states tensor.
N (Tensor | None) – Population size tensor.
n (Tensor | None) – Number of draws tensor.
validate_args (bool | None) – Whether to enable argument validation.

check_inputs(K, N, n, event_shape)[source]¶

Validate hypergeometric parameters.

check_support(data)[source]¶

Check if data is in support of the Hypergeometric distribution.

Delegates to the _HypergeometricDistribution’s check_support method.

Parameters:: data (Tensor) – Input data tensor.
Return type:: Tensor
Returns:: Boolean tensor indicating which values are in support.

params()[source]¶

Returns distribution parameters.

Return type:: dict[str, Tensor]

Non-Parametric Distributions¶

PiecewiseLinear¶

Non-parametric distribution that approximates data using piecewise linear functions from histograms. Can handle both continuous and discrete data.

class spflow.modules.leaves.PiecewiseLinear(scope, out_channels=1, num_repetitions=1, alpha=0.0)[source]¶

Bases: LeafModule

Piecewise linear leaf distribution module.

First constructs histograms from the data using K-means clustering, then approximates the histograms with piecewise linear functions.

This leaf requires initialization with data via the initialize() method before it can be used for inference or sampling.

alpha¶: Laplace smoothing parameter.

xs¶: Nested list of x-coordinates for piecewise linear functions.

ys¶: Nested list of y-coordinates (densities) for piecewise linear functions.

domains¶: List of Domain objects describing each feature.

is_initialized¶: Whether the distribution has been initialized with data.

__init__(scope, out_channels=1, num_repetitions=1, alpha=0.0)[source]¶

Initialize PiecewiseLinear leaf module.

Parameters:

scope (Union[Scope, int, List[int]]) – Variable scope (Scope, int, or list[int]).
out_channels (int) – Number of output channels (clusters via K-means).
num_repetitions (int) – Number of repetitions.
alpha (float) – Laplace smoothing parameter (default 0.0).

initialize(data, domains)[source]¶

Initialize the piecewise linear distribution with data.

Uses K-means clustering to create multiple distributions per leaf, then constructs histograms and approximates them with piecewise linear functions.

Parameters:

data (Tensor) – Training data tensor of shape (N, F) where N is batch size and F is the number of features.
domains (List[Domain]) – List of Domain objects, one per feature.

Raises:

ValueError – If data shape doesn’t match scope.

Return type:

None

log_likelihood(data, cache=None)[source]¶

Compute log-likelihoods for input data.

Parameters:

data (Tensor) – Input data tensor of shape (N, F).
cache (Cache | None) – Optional cache dictionary.

Return type:

Tensor

Returns:

Log-likelihood tensor.

params()[source]¶

Returns the parameters of the distribution.

For PiecewiseLinear, returns xs and ys nested lists.

Return type:: dict

reset()[source]¶

Reset the distribution to uninitialized state.

Return type:: None

sample(num_samples=None, data=None, is_mpe=False, cache=None, sampling_ctx=None)[source]¶

Sample from the piecewise linear distribution.

Parameters:

num_samples (int | None) – Number of samples to generate.
data (Tensor | None) – Optional evidence tensor.
is_mpe (bool) – Perform MPE (mode) instead of sampling.
cache (Cache | None) – Optional cache dictionary.
sampling_ctx (Optional[SamplingContext]) – Optional sampling context.

Return type:

Tensor

Returns:

Sampled data tensor.

property distribution: PiecewiseLinearDist¶

Returns the underlying PiecewiseLinearDist object.

Raises:: ValueError – If the distribution has not been initialized.

property mode: Tensor¶

Return distribution mode.

Returns:: Mode of the distribution.