API

Datasets

Helpers to load the example datasets shipped with the package.

load_data(dataset)

Load dataset for continuous or binary HGF.

Updates functions

Update functions are the heart of probabilistic networks as they shape the propagation of beliefs in the neural hierarchy. The library implements the standard variational updates for value and volatility coupling, as described in Weber et al. (2023).

The updates module contains the update functions used during the belief propagation. Update functions are available through three sub-modules, organized according to their functional roles. We usually dissociate the first updates, triggered top-down (from the leaves to the roots of the networks), that are prediction steps and recover the current state of inference. The second updates are the prediction error, signalling the divergence between the prediction and the new observation (for input nodes), or state (for state nodes). Interleaved with these steps are posterior update steps, where a node receives prediction errors from the child nodes and estimates new statistics.

Posterior updates

Update the sufficient statistics of a state node after receiving prediction errors from children nodes. The prediction errors from all the children below the node should be computed before calling the posterior update step.

Categorical nodes

categorical_state_update(attributes, ...)

Update the categorical input node given an array of binary observations.

Continuous nodes

continuous_node_posterior_update_ehgf(...[, ...])	Update the posterior of a continuous node using the eHGF update.
continuous_node_posterior_update_ehgf_mean_field(...)	Update the posterior of a continuous node using the eHGF update with mean-field.

continuous_node_posterior_update_unbounded(...)	Update the posterior of a continuous node with unbounded quadratic approximation.
posterior_update_unbounded(attributes, ...)	Compute unbounded posterior for a continuous volatility-parent node.

posterior_update_mean_continuous_node(...)	Continuous mean update without mean-field approximation.
posterior_update_mean_continuous_node_mean_field(...)	Continuous mean update with mean-field approximation [R4d094550a653-1].

posterior_update_precision_continuous_node(...)	Update the precision of a state node using the volatility prediction errors.
posterior_update_precision_continuous_node_mean_field(...)	Update the precision of a state node using the volatility prediction errors.
posterior_update_precision_continuous_node_ehgf(...)	Route to the relaxed enhanced-HGF precision or missing-value update.
posterior_update_precision_continuous_node_ehgf_mean_field(...)	Route to the enhanced-HGF precision update (mean-field) or the missing path.
precision_update(attributes, edges, node_idx)	Compute new precision in the case of observed values.
precision_update_mean_field(attributes, ...)	Compute new precision in the case of observed values.
precision_update_ehgf(attributes, edges, ...)	Enhanced-HGF precision update for observed values (relaxed value coupling).
precision_update_ehgf_mean_field(attributes, ...)	Enhanced-HGF precision update for observed values (mean-field value coupling).
precision_update_missing_values(attributes, ...)	Compute new precision in the case of missing observations.

continuous_node_posterior_update(attributes, ...)	Update the posterior of a continuous node using the standard HGF update.
continuous_node_posterior_update_mean_field(...)	Update the posterior of a continuous node using the mean-field standard update.

Exponential family

posterior_update_exponential_family_dynamic(...)

Update the hyperparameters of an ef state node using HGF-implied learning rates.

Prediction steps

Compute the expectation for future observation given the influence of parent nodes. The prediction step are executed for all nodes, top-down, before any observation.

Binary nodes

binary_state_node_prediction(attributes, ...)

Get the new expected mean and precision of a binary state node.

Continuous nodes

predict_mean(attributes, edges, node_idx)	Compute the expected mean of a continuous state node.
predict_precision(attributes, edges, node_idx)	Compute the predicted precisions of a continuous state node (no mean-field).
predict_precision_mean_field(attributes, ...)	Compute the expected precision of a continuous state node (mean-field).
continuous_node_prediction(attributes, ...)	Update the expected mean and expected precision of a continuous node [Rad74659bc93e-1].
continuous_node_prediction_mean_field(...)	Mean-field (v0.2.11) prediction step for a continuous node.

Dirichlet processes

dirichlet_node_prediction(edges, attributes, ...)

Prediction of a Dirichlet process node.

Prediction error steps

Compute the value and volatility prediction errors of a given node. The prediction error can only be computed after the posterior update (or observation) of a given node.

Binary state nodes

binary_state_node_prediction_error(...)	Compute the value prediction errors and predicted precision of a binary node.
binary_finite_state_node_prediction_error(...)	Update the posterior of a binary node given finite precision of the input.

Categorical state nodes

categorical_state_prediction_error(...)

Prediction error from a categorical state node.

Continuous state nodes

continuous_node_value_prediction_error(...)	Compute the value prediction error of a state node.
continuous_node_volatility_prediction_error(...)	Compute the volatility prediction error of a state node.
continuous_node_prediction_error(attributes, ...)	Store prediction errors in an input node.

Dirichlet state nodes

dirichlet_node_prediction_error(edges, ...)	Prediction error and update the child networks of a Dirichlet process node.
update_cluster(operands, edges, node_idx)	Update an existing cluster.
create_cluster(operands, edges, node_idx)	Create a new cluster.
get_candidate(value, sensory_precision, ...)	Find the best cluster candidate given previous clusters and an input value.
likely_cluster_proposal(mean_mu_G0, ...[, ...])	Sample likely new belief distributions given pre-existing clusters.
clusters_likelihood(value, expected_mean, ...)	Likelihood of a parametrised candidate under the new observation.

Exponential family

prediction_error_update_exponential_family_fixed(...)	Update the parameters of an exponential family distribution.
prediction_error_update_exponential_family_dynamic(...)	Pass the expected sufficient statistics to the implied continuous nodes.

Observations

Functions for setting observations on input nodes.

set_observation(attributes, node_idx, ...)

Add observations to the target node by setting the posterior to a given value.

Vectorised updates

Layer-wise vectorised implementations of the HGF update equations for deep predictive coding networks. These functions operate on entire layers instead of individual nodes and are used internally by pyhgf.model.DeepNetwork.

Volatile layers

vectorised_layer_prediction(child_state, ...)	Predict expected mean/precision for all nodes in a volatile-node layer.
vectorised_root_prediction(layer_state, ...)	Predict the precisions of a layer that has no value parent above it.
vectorised_layer_posterior_update(layer, ...)	Update the value-level posterior for all nodes in a parent layer.
vectorised_posterior_update_precision_value_level(...)	Update the precision of the value level for all nodes in a layer.
vectorised_posterior_update_mean_value_level(...)	Update the mean of the value level for all nodes in a layer.
vectorised_layer_value_prediction_error(layer)	Compute the value prediction error for all nodes in a layer.
vectorised_layer_volatility_prediction_error(layer)	Compute the volatility prediction error for all nodes in a layer.
vectorised_layer_volatility_posterior_standard(layer)	Update the volatility level using the standard ordering.
vectorised_layer_volatility_posterior_ehgf(...)	EHGF volatility-level posterior update (mean first, then safe precision).
vectorised_layer_volatility_posterior_unbounded(...)	Unbounded volatility-level posterior update (Lambert W₀ dual-quadratic).
vectorised_layer_prediction_error(layer, params)	Compute prediction errors and apply the volatility-level posterior update.

Binary layers

vectorised_binary_prediction(child_state, ...)	Predict expected mean and precision for a binary state node layer.
vectorised_binary_prediction_error(layer)	Compute prediction errors for a binary state node layer.

Categorical layers

vectorised_categorical_prediction(...[, ...])	Predict a categorical state node layer: one belief per class.
vectorised_categorical_prediction_error(layer)	Compute the prediction error of a categorical state node layer.

Continuous layers

ValueChild(state, weights, coupling_fn, ...)	The value child of a continuous layer, as seen by its posterior update.
VolatilityChild(state, kappa, params)	The volatility child of a continuous layer, as seen by its posterior update.
vectorised_continuous_prediction(...[, ...])	Predict expected mean and precisions for a layer of continuous nodes.
vectorised_continuous_prediction_error(layer)	Compute the prediction errors a continuous layer sends to its parents.
vectorised_continuous_value_prediction_error(layer)	Compute the value prediction error for all nodes in a continuous layer.
vectorised_continuous_volatility_prediction_error(layer)	Compute the volatility prediction error for all nodes in a continuous layer.
vectorised_continuous_posterior_update(layer)	Dispatch the continuous posterior update on the layer's children.
vectorised_continuous_posterior_update_standard(layer)	Apply the standard HGF posterior update: precision first, then the mean.
vectorised_continuous_posterior_update_ehgf(layer)	Enhanced-HGF posterior update: mean first, then the safe precision.
vectorised_continuous_posterior_update_unbounded(...)	Unbounded posterior update for a pure volatility-parent layer.

Vectorised learning

learning_weights_vectorised(parent_state, ...)	Per-layer weight-learning factors for the vectorised deep network.
resolve_synaptic_uncertainty_settings([...])	Validate the weight-belief rule's settings and fill in the defaults.
vectorised_synaptic_uncertainty_update(...)	One weight matrix's belief update: new mean, new accumulated precision.

Model

The main classes used to create probabilistic networks for binary or continuous inputs.

Network	A predictive coding neural network.
DeepNetwork	Deep predictive coding network with vectorised operations.
DeadGradientWarning	Some weight matrices receive essentially no learning signal.
add_continuous_state	Add continuous state node(s) to a network.
add_constant_state	Add constant-state (bias) node(s) to a network.
add_binary_state	Add binary state node(s) to a network.
add_ef_state	Add exponential family state node(s) to a network.
add_categorical_state	Add categorical state node(s) to a network.
add_dp_state	Add a Dirichlet Process node to a network.
get_couplings	Transform coupling parameter into tuple of indexes and strenghts.
update_parameters	Update the default node parameters using keywords args and dictonary.
insert_nodes	Insert a set of parametrised node in a network.

Weight transplant

Converters that build a DeepNetwork reproducing the forward pass of a trained Equinox module, translating its weights into PyHGF’s layer layout.

from_linear(linear[, leaf_kwargs, ...])	Build a two-layer network computing exactly linear(x).
from_feedforward(fc1, fc2[, activation, ...])	Build a three-layer network computing exactly fc2(activation(fc1(x))).
from_embedding(embedding[, leaf_kwargs, ...])	Build a two-layer network reproducing a table lookup from one-hot inputs.
from_conv(kernel[, bias, leaf_kwargs, ...])	Build a shared-kernel network holding a transplanted convolution kernel.

Declarative construction

Build a DeepNetwork from serialisable layer configurations, for configuration-driven experiments and reproducible hyperparameter sweeps.

LayerConfig	Configuration for one layer in a DeepNetwork.
resolve_coupling_fn	Resolve a coupling function name to a callable.

Mixed pipelines

Declare models as a tree of parts mixing learning PyHGF networks with frozen calculations. Each part translates prediction errors backward through the tree, so locally-learning networks and fixed operations compose into one model.

PCModule	Base class for all parts in a mixed pipeline.
DeepNetworkAdapter	A learning part: a PyHGF DeepNetwork in the pipeline.
EquinoxAdapter	A frozen part: a fixed calculation that routes errors but never learns.
PCSequential	A chain of parts: forward in order to predict, in reverse to update.
Residual	The shortcut junction: output = input + branch(input).
linear_adapter	Build a frozen linear map around an eqx.nn.Linear's parameters.
layer_norm_adapter	Build a frozen LayerNorm around an eqx.nn.LayerNorm's parameters.
gelu_adapter	Build a frozen GELU: the error is multiplied by the slope at the cached input.

Convolution

A convolution as a weight-shared linear map over image patches: patch extraction and pooling are frozen parts, and the single learning part is one small DeepNetwork applied identically to every patch.

conv_output_shape	Return the (height, width) a convolution produces, without running it.
im2col_adapter	Build a frozen patch extractor: images in, one row per patch out.
avg_pool_adapter	Build a frozen average-pool: block-mean forward, uniform-split backward.
max_pool_adapter	Build a frozen max-pool: block-max forward, winner-take-all backward.
spatial_reshape_adapter	Build a frozen reshape between patch rows and the spatial layout.
flatten_adapter	Build a frozen flatten, mapping (batch, C, H, W) to (batch, C*H*W).
conv_patch_network	Build a fresh 2-layer DeepNetwork computing one conv kernel's map.
conv_block	Assemble one conv layer as a part tree: im2col, shared kernel, GELU, pool.

Transformer

A GPT-style Transformer assembled from mixed-pipeline parts, where any slot can hold a frozen calculation or a learning PyHGF network.

MultiHeadAttention	Causal multi-head self-attention as a composite pipeline part.
HybridGPT	The full Transformer as a mixed pipeline, from token ids to logits.
hybrid_from_gpt	Assemble a HybridGPT from a reference Equinox GPT.

Fused pipeline execution

Run a part tree as a single compiled program per training step: the forward walk, the output error, and every local learning step are staged into one graph with all part state passed explicitly.

FusedPipeline	Run a part tree, one compiled program per training step.
step_report	Per-part step magnitudes after a training step, for rate calibration.

Plots

Plotting functionalities to visualise parameters trajectories and correlations after observing new data. We are currently fully supporting Graphviz for network structure and Matplotlib for time series visualisation. NetworkX is also available for some functions.

Matplotlib

plot_trajectories(network[, ci, ...])	Plot the trajectories of the nodes' sufficient statistics and surprise.
plot_correlations(network)	Plot the heatmap correlation of the sufficient statistics trajectories.
plot_nodes(network, node_idxs[, ci, ...])	Plot the trajectory of expected sufficient statistics of a set of nodes.
plot_samples(network[, show_total_surprise, ...])	Plot simulation trajectories for nodes.
plot_layers(network[, layers, variables, ...])	Plot layer-wise parameter trajectories of a DeepNetwork.

Graphviz

plot_network(network)	Visualisation of node network using GraphViz.
plot_deep_network(deep_network[, filename, view])	Visualisation of a fully connected deep network using GraphViz.

Networkx

plot_network(network[, figsize, node_size, ...])

Visualisation of node network using NetworkX and pydot layout.

Response

A collection of response functions. A response function is simply a callable taking at least the HGF instance as input after observation and returning surprise.

first_level_gaussian_surprise(hgf[, ...])	Gaussian surprise at the first level of a probabilistic network.
total_gaussian_surprise(hgf[, ...])	Sum of the Gaussian surprise across the probabilistic network.
first_level_binary_surprise(hgf[, ...])	Time series of binary surprises for all binary state nodes.
binary_softmax(hgf[, ...])	Surprise under the binary softmax model.
binary_softmax_inverse_temperature(hgf[, ...])	Surprise from a binary softmax parametrised by the inverse temperature.

Utils

Utilities for manipulating neural networks.

add_edges(attributes, edges[, kind, ...])	Add a value or volatility coupling link between a set of nodes.
add_parent(attributes, edges, index, ...)	Add a new continuous-state parent node to the attributes and edges of a network.
beliefs_propagation(attributes, inputs, ...)	Update the networks parameters after observing new data point(s).
fill_categorical_state_node(network, ...)	Generate a binary network implied by categorical state(-transition) nodes.
get_input_idxs(edges)	List all possible default inputs nodes.
get_update_sequence(network, volatility_updates)	Generate an update sequence from the network's structure.
list_branches(node_idxs, edges[, branch_list])	Return the branch of a network from a given set of root nodes.
remove_node(attributes, edges, index)	Remove a given node from the network.
sample_node_distribution(attributes, edges, ...)	Sample a value from the distribution of an input node.
sample(network, time_steps[, n_predictions, ...])	Sample n_predictions forward in generative mode.
to_pandas(network)	Export the nodes trajectories and surprise as a Pandas data frame.

Weight initialisation

Initialisation strategies for coupling weight matrices in predictive-coding networks.

xavier_init(n_parents, n_children[, seed])	Xavier / Glorot uniform initialisation.
he_init(n_parents, n_children[, seed])	He / Kaiming normal initialisation.
orthogonal_init(n_parents, n_children[, ...])	Orthogonal initialisation.
sparse_init(n_parents, n_children[, ...])	Sparse initialisation.

Sampling

Low-level helpers used when sampling from a network’s generative distribution.

single_sample(rng_key, initial_state, ...)

Perform a single prediction using the provided RNG key.

Vectorised belief propagation

Scan-based belief propagation used internally by pyhgf.model.DeepNetwork.

propagation_step(network, opt_state, inputs, ...)	Single propagation step through the network.
prediction_pass(network, x)	Forward-only sweep through the network.
run_scan(init_carry, inputs, optimiser, ...)	Run jax.lax.scan over the belief-propagation step.
run_continuous_scan(network, ys, time_steps)	Filter a sequence of observations through a continuous network.
prediction_sweep(network, x)	JIT-compiled top-down prediction sweep.
update_sweep(network, y[, time_step])	JIT-compiled bottom-up prediction-error + posterior sweep.
learn_sweep(network, opt_state, optimiser, ...)	JIT-compiled weight-learning phase.
input_prediction_error(network)	JIT-compiled prediction error at the input (top) layer.
sample_step(network, x, y[, learning_kind, ...])	One full local learning step for one sample, as a pure function.
batch_step(network, opt_state, x, y[, ...])	One batch-synchronous learning step over many samples at once.
apply_precision_increments(network, increments)	Add precision increments (see _precision_increments()) to a network.
batched_prediction_pass(network, x)	Forward-only sweep for a batch of samples, compiled once and reused.
batched_prediction_states(network, x)	Batched forward sweep returning the per-sample swept states.

Math

Math functions and probability densities.

MultivariateNormal()	The multivariate normal as an exponential family distribution [R2fcf60062970-1].
Normal()	The univariate normal as an exponential family distribution [Rf25c987ee749-1].
gaussian_predictive_distribution(x, xi, nu)	Density of the Gaussian-predictive distribution.
gaussian_density(x, mean, precision)	Gaussian density as defined by mean and precision.
binary_surprise(x, expected_mean[, clipping])	Surprise at a binary outcome.
gaussian_surprise(x, expected_mean, ...)	Surprise at an outcome under a Gaussian prediction.
dirichlet_kullback_leibler(alpha_1, alpha_2)	Compute the Kullback-Leibler divergence between two Dirichlet distributions.
binary_surprise_finite_precision(value, ...)	Compute the binary surprise with finite precision.
sigmoid_inverse_temperature(x, temperature)	Compute the sigmoid response function with inverse temperature parameter.
parametrised_sigmoid(x, theta, phi)	Compute the sigmoid parametrised by \(\phi\) and \(\theta\).
smoothed_rectangular(x, theta_l[, phi_l, ...])	Compute the smoothed rectangular weighting function \(b\).
lambert_w0	Principal branch of the Lambert W function for z >= 0.

Typing

Named tuples and PyTree containers describing the structure and state of probabilistic networks.

Network structure

AdjacencyLists(node_type, value_parents, ...)	Indexes to a node's value and volatility parents.
UpdateSequence(prediction_steps, update_steps)	Set of update functions to apply to the network.

Vectorised networks

LayerParams([tonic_volatility_vol, ...])	Per-layer static parameters.
LayerState(mean, precision, expected_mean, ...)	Vectorised per-layer state, as an eqx.Module.
Layer(state, params, weights_mean, ...[, ...])	One layer of the vectorised deep network.
LayerStack(state, params, weights_mean, ...)	N identical layers stacked into one PyTree with a leading (N,) axis.
VectorisedNetwork(layers, ...[, ...])	Complete vectorised network state.
stack_layers(layers)	Combine N identical Layer instances into a single LayerStack.