PyJuice

https://tractables.github.io/pyjuice/
Admin
PyJuice is a Python library for Probabilistic Circuits. It aims to provide a easy-to-use environment...
Tokens:43,901
Snippets:360
Trust Score:9.7
Update:4 months ago
Show doc for...
Tokens:
Raw
### Install PyJuice from Source

Source: https://tractables.github.io/pyjuice/_sources/getting-started/installation.rst.txt

Install PyJuice by cloning the repository and using pip in editable mode. This is useful for development or when needing the latest unreleased features.

```bash
git clone git@github.com:Tractables/pyjuice.git
cd pyjuice
pip install -e .
```

--------------------------------

### Install PyJuice using pip

Source: https://tractables.github.io/pyjuice/_sources/getting-started/installation.rst.txt

Use this command for the simplest installation of PyJuice. Ensure pip is available in your environment.

```bash
pip install pyjuice
```

--------------------------------

### Define a PC for merge example

Source: https://tractables.github.io/pyjuice/_sources/getting-started/tutorials/05_common_transformations.rst.txt

Defines a PC with multiple input nodes for the same variable scope, suitable for demonstrating the merge transformation.

```Python
with juice.set_block_size(block_size = 4):

    i00 = juice.inputs(0, num_node_blocks = 8, dist = dists.Categorical(num_cats = 5))
    i01 = juice.inputs(0, num_node_blocks = 8, dist = dists.Categorical(num_cats = 5))
    i10 = juice.inputs(1, num_node_blocks = 8, dist = dists.Categorical(num_cats = 5))
    i11 = juice.inputs(1, num_node_blocks = 8, dist = dists.Categorical(num_cats = 5))

    ms00 = juice.multiply(i00, i10)
    ms01 = juice.multiply(i01, i11)

    ns = juice.summate(ms00, ms01, num_node_blocks = 1, block_size = 1)
```

--------------------------------

### Define a PC for deepcopy example

Source: https://tractables.github.io/pyjuice/_sources/getting-started/tutorials/05_common_transformations.rst.txt

Defines a sample Probabilistic Circuit (PC) using PyJuice for demonstrating deepcopy functionality. Ensure block_size is set.

```Python
with juice.set_block_size(block_size = 4):

    i00 = juice.inputs(0, num_node_blocks = 8, dist = dists.Categorical(num_cats = 5))
    i10 = juice.inputs(1, num_node_blocks = 8, dist = dists.Categorical(num_cats = 5))
    i11 = juice.inputs(1, num_node_blocks = 8, dist = dists.Categorical(num_cats = 5))

    ms0 = juice.multiply(i00, i10)
    ms1 = juice.multiply(i00, i11)

    ns = juice.summate(ms0, ms1, num_node_blocks = 1)

ns.init_parameters()
```

--------------------------------

### CircuitOptimizer Base Class Usage Example

Source: https://tractables.github.io/pyjuice/python-api/optim.html

Illustrates the general training pattern using a CircuitOptimizer base class. This pattern involves accumulating flows into pc.param_flows and then calling opt.step() for the EM update and accumulator reset.

```python
opt = juice.optim.MiniBatchEM(pc, step_size = 0.1, pseudocount = 0.01)
for x in loader:
    lls = pc(x)
    lls.mean().backward()   # accumulates flows into pc.param_flows
    opt.step()              # EM update, then resets the flow accumulator
```

--------------------------------

### Prepare Data for Conditional Probability Query

Source: https://tractables.github.io/pyjuice/_sources/getting-started/tutorials/04_query_pc.rst.txt

Prepares data and a missing mask for computing conditional probabilities. This setup is similar to marginal probability queries but is used to specify the evidence for conditioning.

```Python
data = torch.tensor([[0, 0], [1, 0]]).to(device)
missing_mask = torch.tensor([[False, True], [False, True]]).to(device) # True for variables to be conditioned on/are missing
```

--------------------------------

### Full-Batch EM Training Loop

Source: https://tractables.github.io/pyjuice/_downloads/e598bde225aeff7f7f120dd63eb87317/01_train_pc.ipynb

An example training loop for full-batch EM. It manually initializes parameter flows, accumulates flows over the entire training set, and then performs a full-batch EM step.

```python
for epoch in range(1, 1+1):
    t0 = time.time()

    # Manually zeroling out the flows
    pc.init_param_flows(flows_memory = 0.0)

    train_ll = 0.0
    for batch in train_loader:
        x = batch[0].to(device)

        # We only run the forward and the backward pass, and accumulate the flows throughout the epoch
        lls = pc(x)
        lls.mean().backward()

        train_ll += lls.mean().detach().cpu().numpy().item()

    # Set step size to 1.0 for full-batch EM
    pc.mini_batch_em(step_size = 1.0, pseudocount = 0.01)

    train_ll /= len(train_loader)

    t1 = time.time()
    test_ll = 0.0
    for batch in valid_loader:
        x = batch[0].to(pc.device)
        lls = pc(x)
        test_ll += lls.mean().detach().cpu().numpy().item()
    
    test_ll /= len(valid_loader)
    t2 = time.time()
    print(f"[Epoch {epoch}/{1}][train LL: {train_ll:.2f}; val LL: {test_ll:.2f}].....[train forward+backward+step {t1-t0:.2f}; val forward {t2-t1:.2f}] ")
```

--------------------------------

### Full-Batch EM Training Loop

Source: https://tractables.github.io/pyjuice/getting-started/tutorials/01_train_pc.html

Example training loop for full-batch EM. It manually zeros out parameter flows, accumulates flows over the entire training set, and then performs a single EM step.

```python
for epoch in range(1, 1+1):
    t0 = time.time()

    # Manually zeroling out the flows
    pc.init_param_flows(flows_memory = 0.0)

    train_ll = 0.0
    for batch in train_loader:
        x = batch[0].to(device)

        # We only run the forward and the backward pass, and accumulate the flows throughout the epoch
        lls = pc(x)
        lls.mean().backward()

        train_ll += lls.mean().detach().cpu().numpy().item()

    # Set step size to 1.0 for full-batch EM
    pc.mini_batch_em(step_size = 1.0, pseudocount = 0.01)

    train_ll /= len(train_loader)

    t1 = time.time()
    test_ll = 0.0
    for batch in valid_loader:
        x = batch[0].to(pc.device)
        lls = pc(x)
        test_ll += lls.mean().detach().cpu().numpy().item()

    test_ll /= len(valid_loader)
    t2 = time.time()
    print(f"[Epoch {epoch}/{1}][train LL: {train_ll:.2f}; val LL: {test_ll:.2f}].....[train forward+backward+step {t1-t0:.2f}; val forward {t2-t1:.2f}] ")
```

--------------------------------

### Mini-Batch EM Training Loop

Source: https://tractables.github.io/pyjuice/getting-started/tutorials/01_train_pc.html

Example training loop for mini-batch EM. It iterates through epochs, performs forward and backward passes, accumulates log-likelihoods, and updates parameters using the optimizer.

```python
for epoch in range(1, 350+1):
    t0 = time.time()
    train_ll = 0.0
    for batch in train_loader:
        x = batch[0].to(device)

        # Forward pass
        lls = pc(x)

        # Backward pass accumulates the parameter flows
        lls.mean().backward()

        train_ll += lls.mean().detach().cpu().numpy().item()

        # Perform a mini-batch EM step (consumes the flows, then resets the accumulator)
        optimizer.step()

    train_ll /= len(train_loader)

    t1 = time.time()
    test_ll = 0.0
    for batch in valid_loader:
        x = batch[0].to(pc.device)
        lls = pc(x)
        test_ll += lls.mean().detach().cpu().numpy().item()

    test_ll /= len(valid_loader)
    t2 = time.time()

    print(f"[Epoch {epoch}/{350}][train LL: {train_ll:.2f}; val LL: {test_ll:.2f}].....[train forward+backward+step {t1-t0:.2f}; val forward {t2-t1:.2f}] ")
```

--------------------------------

### pyjuice.queries.sample

Source: https://tractables.github.io/pyjuice/python-api/generated/pyjuice.queries.sample.html

Draw samples from a PC by performing a top-down ancestral sampling pass. For unconditional sampling, num_samples must be specified. For conditional sampling (conditional = True), run a forward pass on the evidence first (e.g., via `marginal()`); the sampler then reuses the cached pc.node_mars and draws one sample per example in that batch, so num_samples is ignored.

```APIDOC
## pyjuice.queries.sample

### Description
Draw samples from a PC by performing a top-down ancestral sampling pass. For unconditional sampling, num_samples must be specified. For conditional sampling (conditional = True), run a forward pass on the evidence first (e.g., via `marginal()`); the sampler then reuses the cached pc.node_mars and draws one sample per example in that batch, so num_samples is ignored.

### Method
`sample(_pc : TensorCircuit_, _num_samples : int | None = None_, _conditional : bool = False_, __sample_input_ns : bool = True_, __do_calibration : bool = False_, _** kwargs_)

### Parameters
#### Path Parameters
* **pc** (_TensorCircuit_) – the input PC
* **num_samples** (_Optional_[__int__]_) – number of samples to draw; required for unconditional sampling, ignored when conditional = True
* **conditional** (_bool_) – whether to sample conditioned on the evidence cached by a preceding forward pass

### Returns
#### Success Response
a tensor of samples of size [num_samples, num_vars]

Return type:
torch.Tensor
```

--------------------------------

### Mini-Batch EM Training Loop

Source: https://tractables.github.io/pyjuice/_downloads/e598bde225aeff7f7f120dd63eb87317/01_train_pc.ipynb

An example training loop for mini-batch EM. It iterates through epochs, performs forward and backward passes, accumulates loss, and updates parameters using the optimizer.

```python
for epoch in range(1, 350+1):
    t0 = time.time()
    train_ll = 0.0
    for batch in train_loader:
        x = batch[0].to(device)

        # Forward pass
        lls = pc(x)

        # Backward pass accumulates the parameter flows
        lls.mean().backward()

        train_ll += lls.mean().detach().cpu().numpy().item()

        # Perform a mini-batch EM step (consumes the flows, then resets the accumulator)
        optimizer.step()

    train_ll /= len(train_loader)

    t1 = time.time()
    test_ll = 0.0
    for batch in valid_loader:
        x = batch[0].to(pc.device)
        lls = pc(x)
        test_ll += lls.mean().detach().cpu().numpy().item()
    
    test_ll /= len(valid_loader)
    t2 = time.time()

    print(f"[Epoch {epoch}/{350}][train LL: {train_ll:.2f}; val LL: {test_ll:.2f}].....[train forward+backward+step {t1-t0:.2f}; val forward {t2-t1:.2f}] ")
```

--------------------------------

### Compute Conditional Probabilities with Soft Evidence

Source: https://tractables.github.io/pyjuice/_sources/getting-started/tutorials/04_query_pc.rst.txt

This example shows how to compute conditional probabilities when soft evidence (probability distributions over values) is provided. The `data` tensor should represent these distributions, and `missing_mask` still indicates which variables are involved in the conditioning.

```Python
data = torch.tensor([[[0.4, 0.6, 0, 0], [0, 0, 0, 0]], [[0.3, 0.7, 0, 0], [0, 0, 0, 0]]]).to(device)
missing_mask = torch.tensor([[False, True], [False, True]]).to(device)

outputs = juice.queries.conditional(
    pc, data = data, missing_mask = missing_mask, target_vars = [1]
)
```

--------------------------------

### sample

Source: https://tractables.github.io/pyjuice/python-api/queries.html

Draw samples from a PC by performing a top-down ancestral sampling pass.

```APIDOC
## sample

### Description
Draw samples from a PC by performing a top-down ancestral sampling pass.

### Method
[Method not specified in source]

### Endpoint
[Endpoint not specified in source]

### Parameters
[Parameters not specified in source]

### Request Example
[Request example not specified in source]

### Response
[Response details not specified in source]
```

--------------------------------

### Generate a PC

Source: https://tractables.github.io/pyjuice/_sources/getting-started/tutorials/04_query_pc.rst.txt

Creates a simple PC with two variables, X1 and X2, and initializes its parameters. This is a prerequisite for querying.

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

ni0 = juice.inputs(0, num_nodes = 2, dist = dists.Categorical(num_cats = 2))
ni1 = juice.inputs(1, num_nodes = 2, dist = dists.Categorical(num_cats = 4))

ms = juice.multiply(ni0, ni1)
ns = juice.summate(ms, num_nodes = 1)

ns.init_parameters()

pc = juice.compile(ns)
```

--------------------------------

### TensorCircuit.zero_param_flows

Source: https://tractables.github.io/pyjuice/python-api/tensorcircuit.html

Resets all accumulated parameter flows to zero. This is typically used to clear the state before starting a new training phase or after certain operations.

```APIDOC
## TensorCircuit.zero_param_flows

### Description
Zero out parameter flows.

### Method
`zero_param_flows()

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
None
```

--------------------------------

### Clone PC with tied parameters

Source: https://tractables.github.io/pyjuice/_sources/getting-started/tutorials/05_common_transformations.rst.txt

Creates a copy of a PC where parameters are tied to the original. Tied parameters remain synchronized during transformations and learning.

```Python
new_ns2 = juice.deepcopy(ns, tie_params = True)
```

--------------------------------

### External Class Methods

Source: https://tractables.github.io/pyjuice/_sources/python-api/generated/pyjuice.nodes.distributions.External.rst.txt

This snippet details the methods available for the External class, including how to get the number of parameter flows and the total number of parameters.

```APIDOC
## External Class

### Description
Represents an external distribution. This class provides access to methods for querying parameter information.

### Methods

#### `num_param_flows()`
Returns the number of parameter flows for the distribution.

#### `num_parameters()`
Returns the total number of parameters for the distribution.

### Attributes

#### `need_meta_parameters`
A boolean indicating whether meta-parameters are needed.
```

--------------------------------

### Initialize MiniBatchEM Optimizer

Source: https://tractables.github.io/pyjuice/_downloads/e598bde225aeff7f7f120dd63eb87317/01_train_pc.ipynb

Initializes the MiniBatchEM optimizer for training the PC. Requires the PC model, a step size, and a pseudocount value.

```python
optimizer = juice.optim.MiniBatchEM(pc, step_size = 0.1, pseudocount = 0.1)
```

--------------------------------

### Define a PC for Merging

Source: https://tractables.github.io/pyjuice/_downloads/3988e2237cc0401749dcb522c104a647/05_common_transformations.ipynb

Sets up a PC with multiple input nodes defined on the same variable scopes, suitable for demonstrating the `juice.merge` operation.

```python
with juice.set_block_size(block_size = 4):

    i00 = juice.inputs(0, num_node_blocks = 8, dist = dists.Categorical(num_cats = 5))
    i01 = juice.inputs(0, num_node_blocks = 8, dist = dists.Categorical(num_cats = 5))
    i10 = juice.inputs(1, num_node_blocks = 8, dist = dists.Categorical(num_cats = 5))
    i11 = juice.inputs(1, num_node_blocks = 8, dist = dists.Categorical(num_cats = 5))
    
    ms00 = juice.multiply(i00, i10)
    ms01 = juice.multiply(i01, i11)

    ns = juice.summate(ms00, ms01, num_node_blocks = 1, block_size = 1)
```

--------------------------------

### query

Source: https://tractables.github.io/pyjuice/python-api/queries.html

A general-purpose entry point for running queries on a PC.

```APIDOC
## query

### Description
A general-purpose entry point for running queries on a PC.

### Method
[Method not specified in source]

### Endpoint
[Endpoint not specified in source]

### Parameters
[Parameters not specified in source]

### Request Example
[Request example not specified in source]

### Response
[Response details not specified in source]
```

--------------------------------

### CircuitOptimizer step() with Overridden Step Size

Source: https://tractables.github.io/pyjuice/python-api/optim.html

Shows how to use the step() method of a CircuitOptimizer, optionally overriding the default step size for a single update. This allows for per-step learning rate adjustments without a separate scheduler.

```python
opt.step(_step_size = 0.05)
```

--------------------------------

### Clone PC with variable remapping

Source: https://tractables.github.io/pyjuice/_sources/getting-started/tutorials/05_common_transformations.rst.txt

Creates a copy of a PC and reassigns its variables using a mapping. This allows the copied PC to operate on a different set of variables.

```Python
var_mapping = {0: 2, 1: 3}
new_ns3 = juice.deepcopy(ns, var_mapping = var_mapping)
```

--------------------------------

### Basic Training Loop with MiniBatchEM

Source: https://tractables.github.io/pyjuice/_sources/python-api/optim.rst.txt

Illustrates a typical training loop using MiniBatchEM. The backward pass accumulates parameter flows, and opt.step() performs the EM update and resets the flow accumulator.

```python
opt = juice.optim.MiniBatchEM(pc, step_size = 0.1, pseudocount = 0.01)
for x in loader:
    lls = pc(x)
    lls.mean().backward()   # accumulates parameter flows
    opt.step()              # EM update, then resets the flow accumulator
```

--------------------------------

### Create Input Nodes

Source: https://tractables.github.io/pyjuice/python-api/generated/pyjuice.merge.html

Initializes input nodes for a probabilistic circuit. Ensure 'num_node_blocks' and 'dists' are defined in your scope.

```python
>>> i00 = inputs(0, num_node_blocks, dists.Categorical(num_cats = 5))
>>> i01 = inputs(0, num_node_blocks, dists.Categorical(num_cats = 5))
>>> i10 = inputs(1, num_node_blocks, dists.Categorical(num_cats = 5))
>>> i11 = inputs(1, num_node_blocks, dists.Categorical(num_cats = 5))
```

--------------------------------

### Construct HMM using pyjuice context manager

Source: https://tractables.github.io/pyjuice/getting-started/tutorials/03_construct_hmm.html

Builds the HMM structure within a context manager that sets the block size for PC nodes. This includes defining input nodes, emission probabilities, and transition probabilities.

```python
with juice.set_block_size(block_size):
    # We begin by defining p(X_{T-1}|Z_{T-1}) for all k = 0...K-1
    ns_input = juice.inputs(seq_length - 1, num_node_blocks = num_node_blocks,
                            dist = dists.Categorical(num_cats = num_emits))

    ns_sum = None
    curr_zs = ns_input
    for var in range(seq_length - 2, -1, -1):
        # The emission probabilities p(X_{var}|Z_{var}=k) for all k = 0...K-1
        curr_xs = ns_input.duplicate(var, tie_params = True)

        # The transition probabilities p(Z_{var+1}|Z_{var})
        if ns_sum is None:
            # Create both the structure and the transition probabilities
            ns = juice.summate(curr_zs, num_node_blocks = num_node_blocks)
            ns_sum = ns
        else:
            # Create only the structure and reuse the transition probabilities from `ns_sum`
            ns = ns_sum.duplicate(curr_zs, tie_params=True)

        curr_zs = juice.multiply(curr_xs, ns)

    # The Initial probabilities p(Z_{0})
    ns = juice.summate(curr_zs, num_node_blocks = 1, block_size = 1)
```

--------------------------------

### ProdNodes Methods

Source: https://tractables.github.io/pyjuice/_sources/python-api/generated/pyjuice.nodes.ProdNodes.rst.txt

This section details the available methods for the ProdNodes class, including initialization, parameter management, and sparsity checks.

```APIDOC
## ProdNodes

### Description
Represents a product of nodes in the pyjuice library. Provides methods for parameter management and sparsity checks.

### Methods

- **__init__**
  Initializes a ProdNodes object.

- **duplicate**
  Creates a duplicate of the ProdNodes object.

- **get_params**
  Retrieves the parameters of the ProdNodes object.

- **init_parameters**
  Initializes the parameters of the ProdNodes object.

- **is_block_sparse**
  Checks if the ProdNodes object is block sparse.

- **is_sparse**
  Checks if the ProdNodes object is sparse.

- **num_parameters**
  Returns the number of parameters in the ProdNodes object.

```

--------------------------------

### Clone PC with tied parameters and variable remapping

Source: https://tractables.github.io/pyjuice/_sources/getting-started/tutorials/05_common_transformations.rst.txt

Combines parameter tying and variable remapping during PC duplication. The copied PC shares parameters with the original and operates on a new set of variables.

```Python
new_ns3 = juice.deepcopy(ns, tie_params = True, var_mapping = var_mapping)
```

--------------------------------

### Construct HMM nodes with context manager

Source: https://tractables.github.io/pyjuice/_sources/getting-started/tutorials/03_construct_hmm.rst.txt

Builds the HMM structure using a context manager to set the block size for all PC nodes. This involves defining input nodes for emissions and iteratively creating sum nodes for transitions.

```Python
with juice.set_block_size(block_size):
    # We begin by defining p(X_{T-1}|Z_{T-1}) for all k = 0...K-1
    ns_input = juice.inputs(seq_length - 1, num_node_blocks = num_node_blocks,
                            dist = dists.Categorical(num_cats = num_emits))

    ns_sum = None
    curr_zs = ns_input
    for var in range(seq_length - 2, -1, -1):
        # The emission probabilities p(X_{var}|Z_{var}=k) for all k = 0...K-1
        curr_xs = ns_input.duplicate(var, tie_params = True)
    
        # The transition probabilities p(Z_{var+1}|Z_{var})
        if ns_sum is None:
            # Create both the structure and the transition probabilities
            ns = juice.summate(curr_zs, num_node_blocks = num_node_blocks)
            ns_sum = ns
        else:
            # Create only the structure and reuse the transition probabilities from `ns_sum`
            ns = ns_sum.duplicate(curr_zs, tie_params=True)

        curr_zs = juice.multiply(curr_xs, ns)
    
    # The Initial probabilities p(Z_{0})
    ns = juice.summate(curr_zs, num_node_blocks = 1, block_size = 1)
```

--------------------------------

### foreach

Source: https://tractables.github.io/pyjuice/python-api/nodes.html

Traverse all nodes of a PC and apply a specific function.

```APIDOC
## foreach

### Description
Traverse all nodes of a PC and apply a specific function.

### Method
`pyjuice.nodes.foreach`
```

--------------------------------

### pyjuice.load()

Source: https://tractables.github.io/pyjuice/python-api/generated/pyjuice.load.html

Loads a Probabilistic Circuit (PC) from a file. The file must have a '.jpc' extension.

```APIDOC
## pyjuice.load()

### Description
Loads a Probabilistic Circuit (PC) from a file. The file should end with ".jpc".

### Method
N/A (Python function)

### Endpoint
N/A (Python function)

### Parameters
#### Path Parameters
- **fname** (str) - Required - The name of the file from which to load the PC.

### Returns
#### Success Response
- **CircuitNodes** - A PC node representing the loaded circuit.

### Example
```python
from pyjuice import load

# Assuming 'my_circuit.jpc' exists
circuit = load('my_circuit.jpc')
```
```

--------------------------------

### Define Product Nodes with Edge IDs

Source: https://tractables.github.io/pyjuice/_sources/getting-started/tutorials/02_construct_simple_pc.rst.txt

Create product nodes that take input nodes as children, specifying connections using edge_ids. Input node vectors must have the same block_size.

```Python
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)
```

--------------------------------

### Using pyjuice.structural_properties Context Manager

Source: https://tractables.github.io/pyjuice/python-api/generated/pyjuice.structural_properties.html

Use this context manager to temporarily allow non-smooth circuits during circuit construction. Ensure necessary imports like Categorical are available.

```python
with pyjuice.structural_properties(allow_nonsmooth = True):
    nis = pyjuice.inputs(var = 0, num_node_blocks = 4, dist = Categorical(num_cats = 20))
    ....
```

--------------------------------

### sample

Source: https://tractables.github.io/pyjuice/_sources/python-api/queries.rst.txt

Generates samples from the probability distribution defined by the TensorCircuit.

```APIDOC
## sample

### Description
Generates samples from the probability distribution defined by the TensorCircuit.

### Method
(Not specified, likely a function call)

### Parameters
(Parameters not explicitly detailed in the source, but would typically include the TensorCircuit and the number of samples desired.)

### Request Example
```python
# Example usage (assuming tc is a compiled TensorCircuit)
# sample(tc, num_samples)
```

### Response
(Response details not explicitly detailed in the source.)
```

--------------------------------

### pyjuice.compile()

Source: https://tractables.github.io/pyjuice/python-api/generated/pyjuice.compile.html

Compiles a PC represented by a DAG into an equivalent torch.nn.Module.

```APIDOC
## pyjuice.compile()

### Description
Compile a PC represented by a DAG into an equivalent torch.nn.Module.

### Parameters
#### Path Parameters
- **ns** (CircuitNodes) - Required - the root node of the PC’s DAG
- **layer_sparsity_tol** (float) - Optional - the maximum allowed fraction for added pseudo edges within every layer (better to set to a small number for sparse/block-sparse PCs). Defaults to 0.5.
- **max_num_partitions** (Optional[int]) - Optional - maximum number of partitions in a layer
- **disable_gpu_compilation** (bool) - Optional - force PyJuice to use CPU compilation. Defaults to False.
- **force_gpu_compilation** (bool) - Optional - force PyJuice to use GPU compilation. Defaults to False.
- **max_tied_ns_per_parflow_block** (int) - Optional - how many groups of tied parameters are allowed to share the same flow/gradient accumulator (higher values -> consumes less GPU memory; lower values -> potentially avoid stalls caused by atomic operations). Defaults to 8.
- **device** (Optional[Union[int, torch.device]]) - Optional - Which GPU do we use for compilation (the default is torch.cuda.current_device)
- **verbose** (bool) - Optional - Whether to display the progress of the compilation. Defaults to True.

### Returns
- **Module** - the compiled PC with type torch.nn.Module
```

--------------------------------

### Setting Block Size During PC Construction

Source: https://tractables.github.io/pyjuice/_sources/getting-started/tutorials/05_common_transformations.rst.txt

Use `juice.set_block_size` within a context manager to define the block size for nodes during PC construction. This is crucial for efficiency with large PCs.

```Python
with juice.set_block_size(block_size = 2):

    ni0 = juice.inputs(0, num_node_blocks = 2, dist = dists.Categorical(num_cats = 2))
    ni1 = juice.inputs(1, num_node_blocks = 2, dist = dists.Categorical(num_cats = 2))
    ni2 = juice.inputs(2, num_node_blocks = 2, dist = dists.Categorical(num_cats = 2))
    ni3 = juice.inputs(3, num_node_blocks = 2, dist = dists.Categorical(num_cats = 2))

    ms1 = juice.multiply(ni0, ni1, edge_ids = torch.tensor([[0, 0], [0, 1], [1, 0], [1, 1]], dtype = torch.long))
    ns1 = juice.summate(ms1, edge_ids = torch.tensor([[0, 0, 0, 0, 1, 1, 1, 1], [0, 1, 2, 3, 0, 1, 2, 3]], dtype = torch.long))

    ms2 = juice.multiply(ni2, ni3, edge_ids = torch.tensor([[0, 0], [1, 1]], dtype = torch.long))
    ns2 = juice.summate(ms2, edge_ids = torch.tensor([[0, 0, 1, 1], [0, 1, 0, 1]], dtype = torch.long))

    ms = juice.multiply(ns1, ns2, edge_ids = torch.tensor([[0, 0], [1, 1]], dtype = torch.long))
    ns = juice.summate(ms, edge_ids = torch.tensor([[0, 0], [0, 1]], dtype = torch.long), block_size = 1)

ns.init_parameters()
```

--------------------------------

### Create an independent copy of a PC

Source: https://tractables.github.io/pyjuice/_sources/getting-started/tutorials/05_common_transformations.rst.txt

Creates a completely independent copy of an existing PC using juice.deepcopy. No parameters are shared between the original and the copy.

```Python
new_ns1 = juice.deepcopy(ns)
```

--------------------------------

### Import necessary packages

Source: https://tractables.github.io/pyjuice/_sources/getting-started/tutorials/03_construct_hmm.rst.txt

Imports the required libraries for building the HMM.

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

--------------------------------

### PC Creation Functions

Source: https://tractables.github.io/pyjuice/_sources/python-api/pyjuice.rst.txt

Functions for creating and initializing Probabilistic Circuits (PCs).

```APIDOC
## pyjuice.inputs

### Description
Creates a new Probabilistic Circuit.

### Method
`pyjuice.inputs()`

### Endpoint
N/A (Python function)

## pyjuice.multiply

### Description
Multiplies two Probabilistic Circuits.

### Method
`pyjuice.multiply(pc1, pc2)`

### Endpoint
N/A (Python function)

## pyjuice.summate

### Description
Sums two Probabilistic Circuits.

### Method
`pyjuice.summate(pc1, pc2)`

### Endpoint
N/A (Python function)

## pyjuice.set_block_size

### Description
Sets the block size for internal operations.

### Method
`pyjuice.set_block_size(size)`

### Endpoint
N/A (Python function)

## pyjuice.structural_properties

### Description
Retrieves structural properties of a Probabilistic Circuit.

### Method
`pyjuice.structural_properties(pc)`

### Endpoint
N/A (Python function)
```

--------------------------------

### Utility Methods

Source: https://tractables.github.io/pyjuice/_sources/python-api/nodes.rst.txt

Details on utility methods for traversing and aggregating information within PyJuice circuits.

```APIDOC
## Utility Methods

### Description

This section describes utility methods that can be applied to nodes or circuits for operations like iteration and aggregation.

### Methods

- **foreach**: Applies a function to each node in a circuit.
- **foldup_aggregate**: Performs a fold-up aggregation over the circuit.
```

--------------------------------

### Construct HMM using PyJuice context manager

Source: https://tractables.github.io/pyjuice/_downloads/a6f4e70e566387d03c9a1773c5f52cdb/03_construct_hmm.ipynb

Builds the HMM structure iteratively within a `set_block_size` context. This includes defining emission and transition probabilities.

```python
with juice.set_block_size(block_size):
    # We begin by defining p(X_{T-1}|Z_{T-1}) for all k = 0...K-1
    ns_input = juice.inputs(seq_length - 1, num_node_blocks = num_node_blocks,
                            dist = dists.Categorical(num_cats = num_emits))
    
    ns_sum = None
    curr_zs = ns_input
    for var in range(seq_length - 2, -1, -1):
        # The emission probabilities p(X_{var}|Z_{var}=k) for all k = 0...K-1
        curr_xs = ns_input.duplicate(var, tie_params = True)
        
        # The transition probabilities p(Z_{var+1}|Z_{var})
        if ns_sum is None:
            # Create both the structure and the transition probabilities
            ns = juice.summate(curr_zs, num_node_blocks = num_node_blocks)
            ns_sum = ns
        else:
            # Create only the structure and reuse the transition probabilities from `ns_sum`
            ns = ns_sum.duplicate(curr_zs, tie_params=True)

        curr_zs = juice.multiply(curr_xs, ns)
        
    # The Initial probabilities p(Z_{0})
    ns = juice.summate(curr_zs, num_node_blocks = 1, block_size = 1)
```

--------------------------------

### InputNodes Class Methods

Source: https://tractables.github.io/pyjuice/_sources/python-api/generated/pyjuice.nodes.InputNodes.rst.txt

This section details the methods available for the InputNodes class, including initialization, parameter manipulation, and duplication.

```APIDOC
## InputNodes Class

### Description
Represents input nodes in the pyjuice library. Provides methods for managing node parameters and creating duplicates.

### Methods

#### `__init__`
Initializes an InputNodes object.

#### `duplicate()`
Creates a duplicate of the current InputNodes object.

#### `get_params()`
Retrieves the parameters of the InputNodes object.

#### `init_parameters()`
Initializes the parameters for the InputNodes object.

#### `num_parameters()`
Returns the number of parameters in the InputNodes object.

#### `set_meta_params()`
Sets the meta-parameters for the InputNodes object.

#### `set_params()`
Sets the parameters for the InputNodes object.
```

--------------------------------

### PC Creation

Source: https://tractables.github.io/pyjuice/python-api/pyjuice.html

Functions for constructing Probabilistic Circuit (PC) nodes.

```APIDOC
## PC Creation

### `inputs`

Construct a vector of input nodes defined on the same variable and have the same input distribution.

### `multiply`

Construct a vector of product nodes given a list of children PCs defined on disjoint sets of variables.

### `summate`

Construct a vector of sum nodes given a list of children PCs defined on the same sets of variables.

### `set_block_size`

Context-manager that sets the block size of PC nodes constructed by inputs, multiply, and summate.

### `structural_properties`

Context-manager that controls the assertions of circuit structural properties, including smoothness and decomposability.
```

--------------------------------

### Define a Probabilistic Circuit (PC)

Source: https://tractables.github.io/pyjuice/_downloads/3988e2237cc0401749dcb522c104a647/05_common_transformations.ipynb

Defines a sample PC structure using input nodes, multiplication, and summation operations. Parameters are initialized after definition.

```python
with juice.set_block_size(block_size = 4):

    i00 = juice.inputs(0, num_node_blocks = 8, dist = dists.Categorical(num_cats = 5))
    i10 = juice.inputs(1, num_node_blocks = 8, dist = dists.Categorical(num_cats = 5))
    i11 = juice.inputs(1, num_node_blocks = 8, dist = dists.Categorical(num_cats = 5))
    
    ms0 = juice.multiply(i00, i10)
    ms1 = juice.multiply(i00, i11)

    ns = juice.summate(ms0, ms1, num_node_blocks = 1)

ns.init_parameters()
```

--------------------------------

### Full-Batch EM Optimizer Initialization and Usage

Source: https://tractables.github.io/pyjuice/python-api/optim.html

Initialize and use the FullBatchEM optimizer for full-batch EM updates. Parameter flows are accumulated over the entire dataset before a single EM M-step.

```python
opt = juice.optim.FullBatchEM(pc, pseudocount = 0.01)
for epoch in range(num_epochs):
    for x in loader:
        lls = pc(x)
        lls.mean().backward()   # flows accumulate over the whole epoch
    opt.step()                  # one exact EM update, then reset

```

--------------------------------

### Anemone Optimizer Initialization and Usage

Source: https://tractables.github.io/pyjuice/python-api/optim.html

Initialize and use the Anemone optimizer, a scaled mini-batch EM with momentum. It accumulates flows over `niters_per_update` minibatches and applies momentum before updating.

```python
opt = juice.optim.Anemone(pc, step_size = 0.4, momentum = 0.9,
                          niters_per_update = 8, ddp = True)
for x in loader:
    lls = pc(x)
    lls.mean().backward()
    opt.step()              # fires the update every `niters_per_update` minibatches

```

--------------------------------

### sample

Source: https://tractables.github.io/pyjuice/_sources/python-api/generated/pyjuice.queries.sample.rst.txt

Samples from a given probabilistic model. This function is part of the pyjuice.queries module and is used for probabilistic inference tasks.

```APIDOC
## sample

### Description
Samples from a given probabilistic model. This function is part of the pyjuice.queries module and is used for probabilistic inference tasks.

### Method
(Not specified in source, likely a Python function call)

### Endpoint
(Not applicable, this is an SDK function)

### Parameters
(Parameters are not explicitly detailed in the provided source. Refer to the pyjuice library documentation for specific parameter information.)

### Request Example
```python
import pyjuice

# Assuming 'model' is a pyjuice probabilistic model object
# and 'num_samples' is the desired number of samples
samples = pyjuice.queries.sample(model, num_samples=100)
```

### Response
(The response format is not explicitly detailed in the provided source. It typically returns a collection of samples from the model.)
```

--------------------------------

### Full-batch EM Training Loop

Source: https://tractables.github.io/pyjuice/_sources/getting-started/tutorials/01_train_pc.rst.txt

This snippet illustrates a training loop for full-batch EM. It initializes flows to zero, accumulates flows over the entire training set, and then performs a single EM step with a fixed step size.

```Python
for epoch in range(1, 1+1):
        t0 = time.time()

        # Manually zeroling out the flows
        pc.init_param_flows(flows_memory = 0.0)

        train_ll = 0.0
        for batch in train_loader:
            x = batch[0].to(device)

            # We only run the forward and the backward pass, and accumulate the flows throughout the epoch
            lls = pc(x)
            lls.mean().backward()

            train_ll += lls.mean().detach().cpu().numpy().item()

        # Set step size to 1.0 for full-batch EM
        pc.mini_batch_em(step_size = 1.0, pseudocount = 0.01)

        train_ll /= len(train_loader)

        t1 = time.time()
        test_ll = 0.0
        for batch in valid_loader:
            x = batch[0].to(pc.device)
            lls = pc(x)
            test_ll += lls.mean().detach().cpu().numpy().item()
    
        test_ll /= len(valid_loader)
        t2 = time.time()
        print(f"[Epoch {epoch}/{1}][train LL: {train_ll:.2f}; val LL: {test_ll:.2f}].....[train forward+backward+step {t1-t0:.2f}; val forward {t2-t1:.2f}] ")
```

--------------------------------

### HMM Constructor

Source: https://tractables.github.io/pyjuice/python-api/generated/pyjuice.structures.HMM.html

Constructs Hidden Markov Models with specified sequence length, latent space size, and emission space size. Optional parameters allow for defining homogeneous/inhomogeneous models, block sizes, and pre-defined transition, emission, and initialization parameters.

```APIDOC
## HMM()

### Description
Constructs Hidden Markov Models.

### Parameters
#### Path Parameters
* **seq_length** (int) - Required - sequence length
* **num_latents** (int) - Required - size of the latent space
* **num_emits** (int) - Required - size of the emission space
* **homogeneous** (bool) - Optional - whether to define a homogeneous (or inhomogeneous) HMM
* **block_size** (Optional[int]) - Optional - block size of the PC
* **alpha** (Optional[torch.Tensor]) - Optional - optional transition parameters of size [num_latents, num_latents]
* **beta** (Optional[torch.Tensor]) - Optional - optional emission parameters of size [num_latents, num_emits]
* **gamma** (Optional[torch.Tensor]) - Optional - optional init parameters of size [num_latents]
* **sum_edge_ids_constructor** (Callable) - Optional - optional helper functions to create special edge patterns (e.g., block-sparse)
```

--------------------------------

### pyjuice.queries.query

Source: https://tractables.github.io/pyjuice/python-api/generated/pyjuice.queries.query.html

A general-purpose entry point for running queries on a PC. It runs a forward pass (and optionally a backward pass) of the PC, optionally with custom input-layer functions, and is the common backend of `marginal()`, `conditional()`, and `sample()`.

```APIDOC
## pyjuice.queries.query

### Description
A general-purpose entry point for running queries on a PC. It runs a forward pass (and optionally a backward pass) of the PC, optionally with custom input-layer functions, and is the common backend of `marginal()`, `conditional()`, and `sample()`.

### Parameters

#### Path Parameters
None

#### Query Parameters
None

#### Request Body
None

### Method Signature
`query(_pc : TensorCircuit_, _inputs : Tensor_, _run_backward : bool = True_, _fw_input_fn : str | Callable | None = None_, _bk_input_fn : str | Callable | None = None_, _fw_output_fn : Callable | None = None_, _** kwargs_)`

### Parameters

*   **pc** (_TensorCircuit_) – the input PC
*   **inputs** (_torch.Tensor_) – input tensor of size [B, num_vars], or a custom shape paired with fw_input_fn
*   **run_backward** (_bool_) – whether to run the backward pass after the forward pass
*   **fw_input_fn** (_Optional_[Union_[str, Callable]]) – an optional custom function (or the name of an input-layer method) for the forward pass of input layers
*   **bk_input_fn** (_Optional_[Union_[str, Callable]]) – an optional custom function (or the name of an input-layer method) for the backward pass of input layers
*   **fw_output_fn** (_Optional_[Callable]) – an optional function applied to pc right after the forward pass; if provided, its return value is returned immediately and no backward pass is run

### Returns

the log-likelihoods from the forward pass (when run_backward is False), the output of fw_output_fn (when provided), or None (after a backward pass)
```

--------------------------------

### Define HMM Variables and Parameters

Source: https://tractables.github.io/pyjuice/_downloads/a6f4e70e566387d03c9a1773c5f52cdb/03_construct_hmm.ipynb

This snippet shows how to define the variables and parameters for a simple Hidden Markov Model (HMM) using PyJuice. It includes defining the state variable, observation variable, transition probabilities, and emission probabilities.

```python
from pyjuice.layer import Input, Affine
from pyjuice.model import Model
from pyjuice.nodes import Node, Variable, Multiply

# Define the number of states and observations
num_states = 2
num_observations = 3

# Define the HMM variables
# State variable (hidden)
state = Variable(num_states, name="state")
# Observation variable (visible)
observation = Variable(num_observations, name="observation")

# Define the parameters
# Transition probabilities (from state i to state j)
trans_prob = Node(num_states, num_states, name="transition_prob")
# Emission probabilities (from state i to observation k)
emission_prob = Node(num_states, num_observations, name="emission_prob")

# Construct the HMM model
# The model will have two layers: one for the state and one for the observation.
# The state variable depends on itself in the previous time step (transition).
# The observation variable depends on the current state (emission).

# Create a model instance
model = Model(name="HMM")

# Add the state variable to the model
model.add_nodes([state])

# Add the observation variable to the model
model.add_nodes([observation])

# Define the dependencies
# The state at time t depends on the state at time t-1 (transition)
# This is represented by a transition matrix that maps previous state to current state.
# For simplicity, we'll assume a single transition matrix for all time steps.
# In a real HMM, this would be a recurrent connection.

# The observation at time t depends on the state at time t (emission)
# This is represented by an emission matrix that maps current state to observation.

# For a basic HMM, we can define the structure as follows:
# Layer 1: State variable
# Layer 2: Observation variable, dependent on the state variable

# Let's define a simple structure for demonstration purposes.
# In a full HMM implementation, you would typically use recurrent layers or time-unrolled structures.

# For this example, we'll create a simplified model where the state and observation
# are directly linked, and the transition is implicitly handled by how you'd use it.
# A more accurate HMM would involve time-dependent nodes or recurrent connections.

# Let's redefine to show a more typical PyJuice structure for a sequence model.
# We'll use layers to represent the time steps.

# Define the layers for the HMM
# Input layer for the state
state_layer = Input(layer_type='recurrent', input_dim=num_states, name='state_layer')
# Affine layer for the transition probabilities
trans_affine = Affine(in_features=num_states, out_features=num_states, name='transition_affine')
# Affine layer for the emission probabilities
emission_affine = Affine(in_features=num_states, out_features=num_observations, name='emission_affine')

# Connect the layers to form the HMM structure
# The state at time t depends on the state at time t-1 via the transition probabilities.
# The observation at time t depends on the state at time t via the emission probabilities.

# This is a conceptual representation. Actual HMM construction in PyJuice often involves
# more complex layer compositions or custom node definitions for sequences.

# For a basic example, let's consider a single time step:

# Define the state node
hidden_state = Variable(num_states, name="hidden_state")

# Define the observation node
visible_observation = Variable(num_observations, name="visible_observation")

# Define the parameters for transitions and emissions
# Transition matrix: P(state_t | state_{t-1})
trans_matrix = Node(num_states, num_states, name="transition_matrix")
# Emission matrix: P(observation_t | state_t)
emission_matrix = Node(num_states, num_observations, name="emission_matrix")

# Create a model
model = Model(name="SimpleHMM")

# Add nodes to the model
model.add_nodes([hidden_state, visible_observation])

# Define the relationships (this part is conceptual and depends on how you build the graph)
# In PyJuice, you'd typically define these relationships using layer compositions or direct node connections.

# Example of defining a relationship (conceptual):
# visible_observation depends on hidden_state
# This would be modeled by emission_matrix

# hidden_state depends on previous hidden_state (recurrent)
# This would be modeled by trans_matrix

# To make this runnable, let's create a minimal graph structure:

# Create a root node for the model
root = Node()

# Add the hidden state node
root.add_child(hidden_state)

# Add the observation node, dependent on the hidden state
# This implies a graphical model structure where observation is a child of hidden_state
# In a real HMM, this dependency is conditional on the state.

# For a simple graphical model representation:
# We can use a Multiply node to represent the joint probability P(state, observation)
# P(state, observation) = P(state) * P(observation | state)

# Let's assume P(state) is defined by an initial state distribution (not explicitly shown here)
# and P(observation | state) is defined by the emission matrix.

# Create a node representing the conditional probability P(observation | state)
# This node will take the hidden_state as input and output probabilities for observations.
# This is often done using an Affine layer or similar.

# For simplicity, let's just define the nodes and parameters without a full graph construction.
# The actual graph construction depends on the specific HMM architecture (e.g., time-unrolled).

print("HMM variables and parameters defined:")
print(f"States: {num_states}")
print(f"Observations: {num_observations}")
print(f"Transition Probability Node: {trans_prob}")
print(f"Emission Probability Node: {emission_prob}")
print(f"Hidden State Variable: {hidden_state}")
print(f"Visible Observation Variable: {visible_observation}")

```
Tokens:
Raw