pyhgf.typing.vectorised.LayerStack#

class pyhgf.typing.vectorised.LayerStack(state, params, weights_mean, coupling_fn, add_constant_input, has_volatility_parent, fully_connected, kind, n_layers, weights_precision_delta=None)[source]#

N identical layers stacked into one PyTree with a leading (N,) axis.

state/params have leading axis N (each field shape goes from (n_nodes,) to (N, n_nodes)). weights_mean goes from (n_child, n_self[+1]) to (N, n_child, n_self[+1]). Slice index 0 is the bottommost slice in the stack (closest to layer 0 of the network); slice N-1 is the topmost.

Validation constraints, enforced at build time:

  • The layer immediately below the stack must have the same node count as the stack width (so weights_mean[0] shape matches).

  • weights_mean[k] for k > 0 is a square (W, W+bias) block connecting slice k (parent) to slice k-1 (child) within the stack.

Parameters:

  • state (pyhgf.typing.vectorised.LayerState) – The stacked per-layer state, each field with a leading (N,) axis.

  • params (pyhgf.typing.vectorised.LayerParams) – The stacked per-layer static parameters, each field with a leading (N,) axis.

  • weights_mean (jax.Array) – The stacked incoming weight matrices, shape (N, n_child, n_self[+1]).

  • coupling_fn (Callable) – The coupling function shared by all stacked layers.

  • add_constant_input (bool) – Whether a constant (bias) input column is appended to the weights.

  • has_volatility_parent (bool) – Whether the layers have a volatility parent.

  • fully_connected (bool) – Whether the incoming weights are fully connected.

  • kind (str) – The kind of layer, one of "volatile", "binary", or "categorical".

  • n_layers (int) – The number of stacked layers N.

  • weights_precision_delta (Array | None)

__init__(state, params, weights_mean, coupling_fn, add_constant_input, has_volatility_parent, fully_connected, kind, n_layers, weights_precision_delta=None)#
Parameters:
Return type:

None

Methods

__init__(state, params, weights_mean, ...[, ...])

Attributes

weights_precision_delta

The stacked weight-belief precisions, shape (N, n_child, n_self[+1]), or None when the stack carries no weight belief.

state

params

weights_mean

coupling_fn

add_constant_input

has_volatility_parent

fully_connected

kind

n_layers