pyhgf.model.fused.FusedPipeline#

class pyhgf.model.fused.FusedPipeline(part, error_fn=None, error_normalisation='active')[source]#

Run a part tree, one compiled program per training step.

Holds every part’s mutable state (the network beliefs and weights, the optimiser states) as one explicit pytree, and compiles a single step function: forward walk → error at the output → backward walk with every local learning step → error at the input. The error is formed inside the program by error_fn, so no intermediate crosses a compilation boundary.

The part objects only declare the model; they are not advanced while the executor runs. Call merge() to write the current state back onto them (e.g. to inspect a network’s layers, or to save it).

The carried state’s buffers are donated to the compiled step, so state holds live arrays that the next step() consumes: anything read from it that must survive a step (a belief to compare before and after, a weight snapshot) should be copied first, e.g. with np.asarray.

Parameters:

  • part – The part tree to execute: any composition of DeepNetworkAdapter, EquinoxAdapter, PCSequential, Residual, and MultiHeadAttention, or a full HybridGPT (whose inputs are then integer token ids).

  • error_fn (Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]]) – How the descent error at the tree’s output is formed from the output and the training target, e.g. lambda out, target: out - target for a squared-error objective (the default), or lambda probs, ids: probs - jax.nn.one_hot(ids, vocab) for a categorical head trained on integer targets. Must be a pure function of arrays.

  • error_normalisation (str) –

    How the batch average is normalised when some rows carry no error.

    Every learning part averages its update over the rows it is handed. "active" (default) divides by the number of rows carrying a non-zero error; "rows" divides by the row count.

__init__(part, error_fn=None, error_normalisation='active')[source]#
Parameters:

Methods

__init__(part[, error_fn, error_normalisation])

merge()

Write the held state back onto the wrapped parts; return the tree.

predict(x)

Run the forward pass only — no error, no learning, no state change.

step(x, target)

Advance the held state by one training step, in one compiled call.