pyjuice.nodes.ParameterizedSumNodes

class pyjuice.nodes.ParameterizedSumNodes(num_node_blocks: int, chs: Sequence, edge_ids: Tensor | Sequence[Tensor] | None = None, params: Tensor | None = None, zero_param_mask: Tensor | None = None, block_size: int = 0, parameterization: SumParameterization | None = None, **kwargs)

A vector of sum nodes whose effective parameters are the shared (EM-trained) parameters modified by per-sample tensors supplied externally at call time. It is created by pyjuice.summate(..., parameterization = <descriptor>).

This is the single node class for every external parameterization; how the external tensors modify the parameters is owned by the parameterization descriptor – an SumParameterization such as LowRank. The arrangement mirrors input nodes, where one InputNodes class carries a dist descriptor, and it means a new mode adds a descriptor rather than a node class.

The external tensors are passed per call and are never owned or EM-trained by the node:

ns = juice.summate(ch, num_node_blocks = 1,
                   parameterization = LowRank(rank = 16))
pc = juice.compile(root_ns)

lls = pc(x, sum_external_params = {ns: (A, B)})          # A [batch, nodes, 16], B [batch, 16, children]
pc.backward(x, sum_external_params_grad = {ns: (dA, dB)})

The shared parameters continue to train through the ordinary EM / gradient path, unchanged.

Note:

at compile time these nodes are grouped into their own layer, keyed by parameterization.get_signature(), so the standard sum-layer kernels used by every other node are never branched on.

Note:

the external tensors are keyed by the ns instance. A tied duplicate needs its own entry; passing the same tensors for several of them shares one per-sample modification across all of them.

Parameters:

  • num_node_blocks (int) – number of node blocks

  • chs (Sequence[CircuitNodes]) – sequence of child nodes

  • edge_ids (Optional[Tensor]) – a matrix of size [2, # edges] - every size-2 column vector [i,j] defines a set of edges that fully connect the ith sum node block and the jth child node block

  • parameterization (SumParameterization) – the external parameterization descriptor

  • block_size (int) – block size

__init__(num_node_blocks: int, chs: Sequence, edge_ids: Tensor | Sequence[Tensor] | None = None, params: Tensor | None = None, zero_param_mask: Tensor | None = None, block_size: int = 0, parameterization: SumParameterization | None = None, **kwargs) → None

Methods

duplicate(*args[, tie_params])

Create a duplication of the current node with the same specification (i.e., number of nodes, block size).

gather_parameters(params)

Update parameters from the current node to the compiled pyjuice.TensorCircuit.

get_params([as_matrix])

Get the sum node parameters.

init_parameters([perturbation, recursive, ...])

Randomly initialize node parameters.

num_parameters([count_input_params, recursive])

set_params(params[, normalize, pseudocount])

Set the sum node parameters.

update_param_flows(param_flows[, ...])

Update parameter flows from pyjuice.TensorCircuit to the current node.

update_parameters(params[, clone])

Update parameters from pyjuice.TensorCircuit to the current node.