Unify graph-energy summation behind a single core helper

[aprotsenko24]

Summary

Across the codebase, "sum the graph's per-node energies" was re-implemented inline in many places (sum/jnp.sum loops over structure.nodes). This PR introduces one reusable core function, fabricpc.core.energy.total_graph_energy(state, structure, *, internal_only), and routes every clean call site through it.

This is a behavior-preserving deduplication — no energy computation is removed. The per-node energy functionals and the actual reductions are unchanged; the duplicated loops are replaced by one shared, tested helper.

The helper

def total_graph_energy(state, structure, *, internal_only):
    energy = jnp.array(0.0)
    for name in structure.nodes:
        if internal_only and structure.nodes[name].node_info.in_degree == 0:
            continue
        energy = energy + jnp.sum(state.nodes[name].energy)
    return energy

• internal_only=True skips terminal input nodes (in_degree == 0) — the PC training objective; False sums all nodes (evaluation).
• Returns the summed energy without batch normalization; callers divide by batch_size when they want a per-sample mean (several callers add an MSE term, compare relatively, or want the raw sum, so a built-in mean would be wrong for them).
• Iterates structure.nodes in insertion order (float summation is order-sensitive).
• Exported from fabricpc.core (import block + __all__).

Call sites routed through the helper

File	Site	Mode
training/train.py	get_graph_param_gradient energy	internal_only=True, /batch_size
training/train.py	eval_step energy	internal_only=False, /batch_size
training/train.py	evaluate_transformer per-device internal energy (inside vmap)	internal_only=True; output-MSE term kept separate
training/train_autoregressive.py	train_step_autoregressive energy	internal_only=True, /batch_size
utils/dashboarding/extractors.py	extract_total_energy	float(... internal_only=True)
utils/dashboarding/inference_tracking.py	train_step_with_history energy	internal_only=False
tests/test_fabricpc.py	energy-decrease assertion	internal_only=True (×2)
tests/test_storkey_hopfield.py	energy-decrease assertion	internal_only=True (×2)

Plus three focused unit tests for the helper in tests/test_energy.py.

Intentionally left inline

• training/train.py evaluate_pcn multi-GPU (pmap) path: each node's energy is reshaped across devices and sliced to drop padded samples before summing. The helper sums the full per-node array, so it cannot reproduce the padding-trim — routing it would wrongly include padded samples.
• Per-node jnp.sum(state.energy) returns inside node forward() methods, and jnp.mean-based diagnostics — these are single-node returns or mean (not sum) reductions, not graph-level aggregations.

Verification

JAX_PLATFORMS=cpu PYTHONPATH=. python -m pytest tests/ -q → 217 passed.

The swaps preserve behavior: same node set (the in_degree > 0 keep is equivalent to the helper's in_degree == 0 skip), same insertion order, and the same /batch_size normalization at each caller.