pyjuice.nodes.distributions.SoftEvidenceCategorical

class pyjuice.nodes.distributions.SoftEvidenceCategorical(num_cats: int, _dual_flow_backward: bool = True, sort_soft_evidence: bool = True)

A class representing a Categorical distribution that allows external soft evidence.

Note:

with top-k soft evidence (i.e. when soft_evidence_cat_ids is supplied) the forward is bound by the params gather addressed through cat_ids, so the candidate ORDER dominates its cost: 4.3x on the CoDD/latent training config, 1.67x on the CoDD decode config, for a ~0.05 ms sort. The forward now sorts by category id itself (sort_soft_evidence, on by default) – see sort_soft_evidence_candidates() for why this is semantically a no-op.

The BACKWARD is deliberately left in the caller’s order, because categorical_evidence_logp_grad is indexed by candidate slot and belongs to the caller. A training loop that wants the backward sorted too should sort once at top-k time with sort_soft_evidence_candidates() and pass the sorted arrays through both passes; that also lets the two passes share one _build_dense_index build.

Parameters:

• num_cats (int) – number of categories

• sort_soft_evidence (bool) – sort the candidate axis by category id inside the forward. On by default; override per call with sort_soft_evidence = False in the forward’s kwargs, or process-wide with PYJUICE_SOFTEVI_SORT=0.

__init__(num_cats: int, _dual_flow_backward: bool = True, sort_soft_evidence: bool = True)

Methods

num_param_flows()	The number of parameter flows per node.
num_parameters()	The number of parameters per node.

Attributes

need_meta_parameters

A flag indicating whether users need to pass in meta-parameters to the constructor of InputNodes.