.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "getting-started/tutorials/03_construct_hmm.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_getting-started_tutorials_03_construct_hmm.py>`
        to download the full example code.

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_getting-started_tutorials_03_construct_hmm.py:


Construct an HMM
================

This tutorial demonstrates how to construct an HMM with :code:`pyjuice.inputs`, :code:`pyjuice.multiply`, and :code:`pyjuice.summate`.

.. GENERATED FROM PYTHON SOURCE LINES 7-10

.. code-block:: Python


    # sphinx_gallery_thumbnail_path = 'imgs/juice.png'


.. GENERATED FROM PYTHON SOURCE LINES 11-12

Let's start by importing the necessary packages.

.. GENERATED FROM PYTHON SOURCE LINES 12-17

.. code-block:: Python


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


.. GENERATED FROM PYTHON SOURCE LINES 18-19

We start with specifying the structural parameters of the HMM

.. GENERATED FROM PYTHON SOURCE LINES 19-24

.. code-block:: Python


    seq_length = 32
    num_latents = 2048
    num_emits = 4023


.. GENERATED FROM PYTHON SOURCE LINES 25-27

An important parameter to be determined is the block size, which is crucial for PyJuice to compile efficient models.
Specifically, we want the block size to be large enough so that PyJuice can leverage block-based parallelization.

.. GENERATED FROM PYTHON SOURCE LINES 27-30

.. code-block:: Python


    block_size = min(juice.utils.util.max_cdf_power_of_2(num_latents), 1024)


.. GENERATED FROM PYTHON SOURCE LINES 31-32

The number of node blocks is derived accordingly

.. GENERATED FROM PYTHON SOURCE LINES 32-35

.. code-block:: Python


    num_node_blocks = num_latents // block_size


.. GENERATED FROM PYTHON SOURCE LINES 36-38

We use the context manager `set_block_size` to set the block size of all PC nodes.
In the following we assume `T = seq_length` and `K = num_latents`

.. GENERATED FROM PYTHON SOURCE LINES 38-64

.. code-block:: 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)


.. GENERATED FROM PYTHON SOURCE LINES 65-71

Note that :code:`ns.duplication` is a handy function to create duplications of existing node vectors. 
For input nodes (e.g., :code:`ns_input.duplicate(var, tie_params = True)`) we can specify it to define on a new variable. 
The argument :code:`tie_params = True` means we want to use the same set of parameters in the original and the duplicated node vector.
The parameters will remain tied after parameter learning, structural transformation, etc.

For sum node vectors, :code:`ns.duplicate` allows us to specify a new list of children. However, the child nodes must have the same size (same number of node blocks and block size).


.. _sphx_glr_download_getting-started_tutorials_03_construct_hmm.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: 03_construct_hmm.ipynb <03_construct_hmm.ipynb>`

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: 03_construct_hmm.py <03_construct_hmm.py>`

    .. container:: sphx-glr-download sphx-glr-download-zip

      :download:`Download zipped: 03_construct_hmm.zip <03_construct_hmm.zip>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_