mixle.stats.compute.gradient module

Gradient-fitting state objects owned by distribution-layer hooks.

exception GradientFitError

Bases: NotImplementedError

Raised when generic gradient fitting cannot handle a model family.

prior_zero(torch, engine, ref=None)

Return a scalar zero on the active Torch/device context.

prior_family(prior)

Return the normalized prior family name, if prior is mapping-like.

prior_sequence(values, n)

Pad or truncate a sequence of child priors to length n.

Parameters:

n (int)

composite_child_priors(priors, n)

Return child priors for a composite-style product distribution.

Parameters:

n (int)

conditional_priors(priors, keys)

Return per-condition, default, and given priors for conditional models.

record_child_priors(priors, fields, n)

Return field-aligned child priors for a named record model.

Parameters:

n (int)

transform_prior(priors)

Return the base-child prior for a transform wrapper.

select_child_priors(priors, n)

Return choice-child priors for a select/routed model.

Parameters:

n (int)

mixture_priors(priors, n)

Return component priors plus an optional weight prior for mixtures.

Parameters:

n (int)

sequence_priors(priors)

Return element and length priors for an iid sequence model.

markov_chain_priors(priors, row_keys)

Return initial, transition-row, and length priors for Markov chains.

dirichlet_alpha_tensor(alpha, labels, logits, engine, torch)

Broadcast Dirichlet concentration values to a logits tensor.

normal_gamma_log_prior(mu, sigma2, priors, torch)

Return a Normal-Gamma log prior over a Gaussian-style mean/variance.

class CategoricalGradientFitState(template, labels, logits)

Bases: object

Autograd state for finite categorical simplex maps.

Parameters:

• template (Any)

• labels (Sequence[Any])

• logits (Any)

shadow(torch, shadow_child)

Build a temporary distribution object backed by live Torch logits.

score(enc, engine, torch, score_child)

Score encoded categorical observations through the shadow object.

build(torch, build_child, detach_value)

Reconstruct a fitted categorical distribution from optimized logits.

log_prior(priors, prior_strength, torch, engine, initial_leaves_by_id, prior_child)

Return the Dirichlet or weak fallback prior contribution.

Parameters:

prior_strength (float)

class OptionalGradientFitState(template, child, logit_p)

Bases: object

Autograd state for optional/missing-value wrappers.

Parameters:

• template (Any)

• child (Any)

• logit_p (Any)

shadow(torch, shadow_child)

Build a temporary optional wrapper backed by live child/raw p state.

score(enc, engine, torch, score_child)

Score encoded optional observations, including missing values.

build(torch, build_child, detach_value)

Reconstruct a fitted optional distribution from optimized state.

log_prior(priors, prior_strength, torch, engine, initial_leaves_by_id, prior_child)

Return observed-child plus optional missingness prior contribution.

Parameters:

prior_strength (float)

class CompositeGradientFitState(template, children)

Bases: object

Autograd state for product distributions over tuple-like observations.

Parameters:

• template (Any)

• children (Sequence[Any])

shadow(torch, shadow_child)

Build a temporary composite with live child shadows.

score(enc, engine, torch, score_child)

Return summed child log densities for encoded tuple fields.

build(torch, build_child, detach_value)

Reconstruct a fitted composite from fitted children.

log_prior(priors, prior_strength, torch, engine, initial_leaves_by_id, prior_child)

Return the sum of field-aligned child prior contributions.

Parameters:

prior_strength (float)

class ConditionalGradientFitState(template, dmap, default_child, given_child)

Bases: object

Autograd state for conditional keyed children.

Parameters:

• template (Any)

• dmap (Mapping[Any, Any])

• default_child (Any)

• given_child (Any)

shadow(torch, shadow_child)

Build a temporary conditional model with live child shadows.

score(enc, engine, torch, score_child)

Score grouped conditional encodings by condition key.

build(torch, build_child, detach_value)

Reconstruct a fitted conditional distribution from children.

log_prior(priors, prior_strength, torch, engine, initial_leaves_by_id, prior_child)

Return keyed/default/given prior contributions for all active children.

Parameters:

prior_strength (float)

class RecordGradientFitState(template, children)

Bases: CompositeGradientFitState

Autograd state for named record distributions.

Parameters:

• template (Any)

• children (Sequence[Any])

score(enc, engine, torch, score_child)

Score named-record encodings using product-distribution logic.

build(torch, build_child, detach_value)

Reconstruct a fitted record distribution preserving field sources.

log_prior(priors, prior_strength, torch, engine, initial_leaves_by_id, prior_child)

Return field-aligned child prior contributions.

Parameters:

prior_strength (float)

class SelectGradientFitState(template, children)

Bases: CompositeGradientFitState

Autograd state for choice-routed child distributions.

Parameters:

• template (Any)

• children (Sequence[Any])

score(enc, engine, torch, score_child)

Score each routed subset with the selected child model.

build(torch, build_child, detach_value)

Reconstruct a fitted select distribution preserving the router.

log_prior(priors, prior_strength, torch, engine, initial_leaves_by_id, prior_child)

Return child prior contributions aligned to select choices.

Parameters:

prior_strength (float)

class TransformGradientFitState(template, child)

Bases: object

Autograd state for fixed transforms with differentiable children.

Parameters:

• template (Any)

• child (Any)

shadow(torch, shadow_child)

Build a temporary transform wrapper with a live base child.

score(enc, engine, torch, score_child)

Score transformed encodings with optional Jacobian correction.

build(torch, build_child, detach_value)

Reconstruct a transform wrapper around the fitted base child.

log_prior(priors, prior_strength, torch, engine, initial_leaves_by_id, prior_child)

Return the base-child prior contribution for a transform wrapper.

Parameters:

prior_strength (float)

class SequenceGradientFitState(template, child, len_child)

Bases: object

Autograd state for iid sequence distributions.

Parameters:

• template (Any)

• child (Any)

• len_child (Any)

shadow(torch, shadow_child)

Build a temporary sequence model with live element/length children.

score(enc, engine, torch, score_child)

Score flattened sequence elements plus optional length model.

build(torch, build_child, detach_value)

Reconstruct a fitted sequence distribution from child fits.

log_prior(priors, prior_strength, torch, engine, initial_leaves_by_id, prior_child)

Return element and length child prior contributions.

Parameters:

prior_strength (float)

class MixtureGradientFitState(template, components, w_logits)

Bases: object

Autograd state for mixture components and simplex weights.

Parameters:

• template (Any)

• components (Sequence[Any])

• w_logits (Any)

shadow(torch, shadow_child)

Build a temporary mixture backed by live component/weight tensors.

score(enc, engine, torch, score_child)

Return mixture log densities via log-sum-exp over components.

build(torch, build_child, detach_value)

Reconstruct a fitted mixture from optimized components and weights.

log_prior(priors, prior_strength, torch, engine, initial_leaves_by_id, prior_child)

Return component prior contributions plus mixture-weight prior.

Parameters:

prior_strength (float)