HiddenParamOpRelation

HiddenParamOpRelation#

class braintrace.HiddenParamOpRelation#

Connection between an ETP primitive, its trainable parameters, and hidden states.

Records the structural relationship

\[h^t = f(y), \quad y = \mathrm{primitive}(x, \theta)\]

discovered by the compiler for a single ETP primitive equation.

Variables:

  • primitive (Primitive) – The JAX primitive (etp_mm_p, etp_mv_p, etc.).

  • x_var (Var or None) – Jaxpr Var for the input (None for element-wise ops).

  • y_var (Var) – Jaxpr Var for the primitive output.

  • hidden_groups (list of HiddenGroup) – Hidden groups that this op feeds into.

  • y_to_hidden_group_jaxprs (list of Jaxpr) – Transition jaxpr from y to each hidden group.

  • connected_hidden_paths (list of Path) – Hidden-state paths connected to this op.

  • eqn_params (dict) – Static parameters of the primitive equation.

  • path_classification (dict) – Mapping {hidden_path: PathClassification.*} for each connected hidden state. Populated by the path-classification pass.

  • trainable_vars (dict) – Per-key dict mapping a primitive-chosen key name (e.g. 'weight', 'bias', 'lora_b', 'lora_a') to its jaxpr Var, with one entry per declared trainable input.

  • trainable_paths (dict) – Per-key dict mapping each key to the owning ParamState’s module path. When two keys trace to the same ParamState (e.g. a merged {weight, bias} Linear), the entries share a path.

  • trainable_leaf_indices (dict) – Per-key dict mapping each key to the leaf index in jax.tree.leaves of the owning ParamState.

  • trainable_param_states (dict) – Per-key dict mapping each key to the actual ParamState object.

  • trainable_processing_chains (dict) – Per-key dict mapping each key to the backward-trace processing chain (primitives traversed from the trainable invar back to the originating ParamState invar).

  • control_flow_context (RelationDescent or None) – Descent context when this relation lives inside a descended scan body (Phase 4 structured scan descent); None for flat relations.

Notes

The dict-typed fields default to class-level empty dicts that are SHARED across instances (the usual NamedTuple default semantics). The compiler always passes freshly-built dicts and no consumer mutates a relation, so this never bites in the pipeline — but when hand-constructing relations in tests, pass explicit dicts instead of mutating the defaults.

static __new__(_cls, primitive, x_var, y_var, hidden_groups, y_to_hidden_group_jaxprs, connected_hidden_paths, eqn_params, path_classification={}, trainable_vars={}, trainable_paths={}, trainable_leaf_indices={}, trainable_param_states={}, trainable_processing_chains={}, control_flow_context=None)#

Create new instance of HiddenParamOpRelation(primitive, x_var, y_var, hidden_groups, y_to_hidden_group_jaxprs, connected_hidden_paths, eqn_params, path_classification, trainable_vars, trainable_paths, trainable_leaf_indices, trainable_param_states, trainable_processing_chains, control_flow_context)

connected_hidden_paths#

Alias for field number 5

control_flow_context#

Alias for field number 13

dict()#

Return this relation’s named fields as a plain dictionary.

Returns:

dict – An ordered mapping from field name to value, as produced by the underlying typing.NamedTuple.

eqn_params#

Alias for field number 6

hidden_groups#

Alias for field number 3

path_classification#

Alias for field number 7

primitive#

Alias for field number 0

trainable_leaf_indices#

Alias for field number 10

trainable_param_states#

Alias for field number 11

trainable_paths#

Alias for field number 9

trainable_processing_chains#

Alias for field number 12

trainable_vars#

Alias for field number 8

x_var#

Alias for field number 1

y_to_hidden_group_jaxprs#

Alias for field number 4

y_to_hidden_groups(y_val, const_vals, concat_hidden_vals=True)#

Evaluate the transition jaxprs mapping y to hidden-group values.

Parameters:

  • y_val (jax.Array) – The value of the primitive output y.

  • const_vals (dict) – Mapping from each transition-jaxpr constvar to its value.

  • concat_hidden_vals (bool, optional) – If True, concatenate each group’s hidden values into a single array via HiddenGroup.concat_hidden(). Default True.

Returns:

list – One entry per hidden group: either a list of per-state arrays (when concat_hidden_vals is False) or a single concatenated array (when True).

y_var#

Alias for field number 2

HiddenParamOpRelation.__init__()#