Construct Simple PCs

In this tutorial, you will learn about the basic APIs to construct PCs.

# sphinx_gallery_thumbnail_path = 'imgs/juice.png'

Let’s start by importing the necessary packages.

import torch
import pyjuice as juice
import pyjuice.nodes.distributions as dists

A PC can be created using pyjuice.inputs, pyjuice.multiply, and pyjuice.summate, which create input node vectors, product node vectors, and sum node vectors, respectively. The goal of this tutorial is to get you familiar with these functions.

Input nodes

Let us start with pyjuice.inputs.

ni0 = juice.inputs(var = 0, num_node_blocks = 4, block_size = 2, dist = dists.Categorical(num_cats = 4))

The above line defines num_node_blocks * block_size = 8 input nodes on variable ID var = 0 featuring Categorical distributions with 4 categories (i.e., dists.Categorical(num_cats = 4)). The set of input distributions are defined under pyjuice.nodes.distributions.

A seemingly redundant pair of parameters are num_node_blocks and block_size, since we can alternatively only specify the number of nodes in the node vector ni0. In fact, ensuring large block_size is crucial to the efficiency of the PC. So always try to use large block sizes when defining PCs. Although using larger block sizes when defining input node vectors do not have any negative effects, it poses restrictions on the edge connection pattern as we shall proceed to show.

Note that block_size has to be a power of 2.

Product nodes

Let us define another input node vector and a product node vector that take the input nodes as children.

ni1 = juice.inputs(var = 1, num_node_blocks = 4, block_size = 2, dist = dists.Categorical(num_cats = 4))

edge_ids = torch.tensor([[0, 0], [1, 1], [2, 2], [3, 3]])
ms = juice.multiply(ni0, ni1, edge_ids = edge_ids)

ms defines a product node vector where every node have two children: the first child is a node in ni0 and the second child is a node in ni1. edge_ids specifies which child nodes do every node in ms connects to. Specifically, edge_ids has size [# product node blocks, # child node vectors], so in this case, it should has size [4, 2]. Specifically, edge_ids[i,j] suggests that the i th product node block connects to the edge_ids[i,j] th node block in the j th child node vector (assume we always count from 0). For example, edge_ids[1,0] = 1 means that the 1th product node block connects to the 1th node block in ni0.

We require the node vectors fed to pyjuice.multiply have the same block_size. And the block size of the output product node vector is also the same with that of the inputs. We do not need to specify the number of node blocks (e.g., using num_node_blocks) since it is equal to edge_ids.size(0). For pyjuice.multiply, if two node blocks are connected (as defined by edge_ids), we assume the :math:`i`th node in the (parent) product node block is connected to the :math:`i`th node in the child node block.

When we do not provide the edge_ids, PyJuice assumes it to be the following (we can only use this shortcut when the child node vectors have the same num_node_blocks and block_size):

num_node_blocks = 4
num_child_node_vectors = 2
edge_ids = torch.arange(0, num_node_blocks)[:,None].repeat(1, num_child_node_vectors)

Sum nodes

Finally, we introduce pyjuice.summate, which is used to define sum node vectors.

edge_ids = torch.tensor([[0, 0, 1, 1, 2, 2, 3, 3, 4, 5], [0, 1, 0, 2, 1, 2, 2, 3, 2, 0]])
ns = juice.summate(ms, num_node_blocks = 6, block_size = 2, edge_ids = edge_ids)

Similar to pyjuice.multiply, the positional arguments to pyjuice.summate are the list of input node vectors (here we only provide ms). num_node_blocks and block_size are the number of node blocks and the size of each node block, respectively. Therefore, ns defines a vector of num_node_blocks * block_size = 12 sum nodes that are (partially) connected to the nodes in ms. If there are multiple child node vectors, pyjuice.summate assumes they have the same block size. However, the sum node vector can have a different block size compared to its children.

The connection pattern is specified by the keyword argument edge_ids, which have shape [2, # edge blocks]. Every size-2 column vector \([m, n]^T\) in edge_ids indicates the existance of fully-connected edges between the m th sum node block and the n th product node block. That is, in the case where both ns and ms have block size 2, every column in edge_ids specifies \(2 \times 2 = 4\) edges.

If edge_ids is not provided, we assume that all nodes in the sum node vector are connected to all child nodes:

ns = juice.summate(ms, num_node_blocks = 6, block_size = 2)