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=1, num_repetitions=1, parameter_fn=None, validate_args=True, loc=None, scale=None)[source]¶
Bases:
LeafModuleNormal (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=1, num_repetitions=1, parameter_fn=None, validate_args=True, loc=None, scale=None)[source]¶
Initialize Normal distribution.
- Parameters:
scope – Variable scope. Can be a Scope object, a single integer, or an iterable of integers (list, tuple, numpy array, torch tensor, etc.).
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.
LogNormal¶
Log-normal distribution (exponential of a normal distribution).
- class spflow.modules.leaves.LogNormal(scope, out_channels=1, num_repetitions=1, parameter_fn=None, validate_args=True, loc=None, scale=None)[source]¶
Bases:
LeafModuleLog-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=1, 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.
Exponential¶
Exponential distribution with learnable rate parameter.
- class spflow.modules.leaves.Exponential(scope, out_channels=1, num_repetitions=1, parameter_fn=None, validate_args=True, rate=None)[source]¶
Bases:
LeafModuleExponential 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=1, 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.
Laplace¶
Laplace distribution with learnable location and scale parameters.
- class spflow.modules.leaves.Laplace(scope, out_channels=1, num_repetitions=1, parameter_fn=None, validate_args=True, loc=None, scale=None)[source]¶
Bases:
LeafModuleNormal (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=1, 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.
Gamma¶
Gamma distribution with learnable alpha and beta parameters.
- class spflow.modules.leaves.Gamma(scope, out_channels=1, num_repetitions=1, parameter_fn=None, validate_args=True, concentration=None, rate=None)[source]¶
Bases:
LeafModuleGamma 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=1, 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.
Uniform¶
Uniform distribution over a bounded interval.
- class spflow.modules.leaves.Uniform(scope, out_channels=1, num_repetitions=1, low=None, high=None, validate_args=True, support_outside=True)[source]¶
Bases:
LeafModuleUniform 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=1, 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) – Number of output channels (inferred from params if None).num_repetitions (
int) – Number of repetitions.validate_args (
bool|None) – Whether to enable torch.distributions argument validation.support_outside (
bool) – Whether values outside [start, end] are supported.
Discrete Distributions¶
Categorical¶
Categorical distribution over K categories with learnable probabilities.
- class spflow.modules.leaves.Categorical(scope, out_channels=1, num_repetitions=1, K=None, probs=None, logits=None, parameter_fn=None, validate_args=True)[source]¶
Bases:
LeafModuleCategorical 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=1, 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) – 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).parameter_fn (
Module|None) – Optional neural network for parameter generation.validate_args (
bool|None) – Whether to enable torch.distributions argument validation.
Bernoulli¶
Bernoulli distribution (binary outcomes) with learnable probability.
- class spflow.modules.leaves.Bernoulli(scope, out_channels=1, num_repetitions=1, parameter_fn=None, validate_args=True, probs=None, logits=None)[source]¶
Bases:
LeafModuleBernoulli 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=1, 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.
Binomial¶
Binomial distribution with learnable probability and fixed number of trials.
- class spflow.modules.leaves.Binomial(scope, out_channels=1, num_repetitions=1, total_count=None, probs=None, logits=None, parameter_fn=None, validate_args=True, differentiable_sampling_method=BinomialDifferentiableSamplingMethod.NORMAL)[source]¶
Bases:
LeafModuleBinomial 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=1, num_repetitions=1, total_count=None, probs=None, logits=None, parameter_fn=None, validate_args=True, differentiable_sampling_method=BinomialDifferentiableSamplingMethod.NORMAL)[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.differentiable_sampling_method (
BinomialDifferentiableSamplingMethod) – Method for differentiable sampling (default: NORMAL).
Poisson¶
Poisson distribution with learnable rate parameter.
- class spflow.modules.leaves.Poisson(scope, out_channels=1, num_repetitions=1, parameter_fn=None, validate_args=True, rate=None)[source]¶
Bases:
LeafModulePoisson 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=1, 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.
Geometric¶
Geometric distribution with learnable success probability.
- class spflow.modules.leaves.Geometric(scope, out_channels=1, num_repetitions=1, parameter_fn=None, validate_args=True, probs=None, logits=None)[source]¶
Bases:
LeafModuleGeometric 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=1, 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 (
int) – 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.logits (
Tensor|None) – Log-odds tensor of the success probability.
NegativeBinomial¶
Negative binomial distribution with learnable parameters.
- class spflow.modules.leaves.NegativeBinomial(scope, out_channels=1, num_repetitions=1, total_count=None, probs=None, logits=None, parameter_fn=None, validate_args=True)[source]¶
Bases:
LeafModuleNegative Binomial distribution leaf matching
torch.distributions.NegativeBinomial.In PyTorch,
NegativeBinomial(total_count=r, probs=p)models the number of successes observed beforerfailures occur, where each trial succeeds with probabilityp. This leaf uses that exact parameterization.Notes
total_count(r) is fixed and cannot be learned.probs(p) is learnable and stored in logit-space for numerical stability.
- total_count¶
Fixed number of failures before stopping (buffer).
- probs¶
Success probability in
[0, 1](stored in logit-space).
- distribution¶
Underlying
torch.distributions.NegativeBinomial.
- __init__(scope, out_channels=1, 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 failures before stopping (required).probs (
Tensor|None) – Success probability tensor (optional).parameter_fn (
Module) – Optional neural network for parameter generation.validate_args (
bool|None) – Whether to enable torch.distributions argument validation.
Hypergeometric¶
Hypergeometric distribution.
- class spflow.modules.leaves.Hypergeometric(scope, out_channels=1, num_repetitions=1, K=None, N=None, n=None, validate_args=True)[source]¶
Bases:
LeafModuleHypergeometric 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=1, num_repetitions=1, K=None, N=None, n=None, validate_args=True)[source]¶
Initialize Hypergeometric distribution leaf module.
CLTree¶
Chow–Liu tree multivariate discrete leaf with exact inference under missing values.
- class spflow.modules.leaves.CLTree(scope, out_channels=1, num_repetitions=1, *, K, alpha=0.01, parents=None, log_cpt=None, validate_args=True)[source]¶
Bases:
LeafModuleChow-Liu tree discrete multivariate leaf.
Notes
log_likelihood() supports NaNs by exact marginalization (belief propagation).
maximum_likelihood_estimation() learns structure once (if missing) and then updates CPTs. Structure is shared across channels/repetitions.
sample() supports conditional sampling and MPE completion under evidence.
- fit_structure(data, *, sample_weights=None)[source]¶
Learn Chow-Liu structure from complete discrete data.
Non-Parametric Distributions¶
Histogram¶
Piecewise-constant histogram density over fixed bin edges with learnable bin probabilities.
- class spflow.modules.leaves.Histogram(scope, *, bin_edges, out_channels=1, num_repetitions=1, probs=None, logits=None, min_prob=1e-12, validate_args=True)[source]¶
Bases:
LeafModuleHistogram leaf distribution (continuous piecewise-constant density).
The histogram uses fixed bin edges and learnable bin probabilities (stored as logits). Within each bin, the density is constant:
p(bin) / width(bin).Notes
This leaf is univariate:
len(scope.query) == 1.Values outside the support
[bin_edges[0], bin_edges[-1])have log-likelihood-inf.NaN values are marginalized out by
log_likelihood()and contribute0to the log-likelihood.MPE returns the midpoint of the maximum-probability bin (per channel/repetition).
- distribution(with_differentiable_sampling=False)[source]¶
Return this leaf’s distribution.
- Parameters:
with_differentiable_sampling (
bool) – Hook for subclasses to return an alternative differentiable distribution when sampling requires gradient flow. Ignored by the base implementation.- Return type:
HistogramDist
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:
LeafModulePiecewise 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.
- distribution(with_differentiable_sampling=False)[source]¶
Return the underlying PiecewiseLinearDist object.
- Parameters:
with_differentiable_sampling (
bool) – Whether to request a differentiable sampling distribution.- Raises:
ValueError – If the distribution has not been initialized.
- Return type:
PiecewiseLinearDist
- 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:
- Raises:
ValueError – If data shape doesn’t match scope.
- Return type: