radu
/
SSMBook
镜像来自 https://github.com/probml/ssm-book.git


			
							123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867
							
<!DOCTYPE html>

<html>
  <head>
    <meta charset="utf-8" />
    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
    <title>HMM filtering (forwards algorithm) &#8212; State Space Models: A Modern Approach</title>
    
  <link href="../../_static/css/theme.css" rel="stylesheet">
  <link href="../../_static/css/index.ff1ffe594081f20da1ef19478df9384b.css" rel="stylesheet">

    
  <link rel="stylesheet"
    href="../../_static/vendor/fontawesome/5.13.0/css/all.min.css">
  <link rel="preload" as="font" type="font/woff2" crossorigin
    href="../../_static/vendor/fontawesome/5.13.0/webfonts/fa-solid-900.woff2">
  <link rel="preload" as="font" type="font/woff2" crossorigin
    href="../../_static/vendor/fontawesome/5.13.0/webfonts/fa-brands-400.woff2">

    
    <link rel="stylesheet" type="text/css" href="../../_static/pygments.css" />
    <link rel="stylesheet" type="text/css" href="../../_static/sphinx-book-theme.css?digest=c3fdc42140077d1ad13ad2f1588a4309" />
    <link rel="stylesheet" type="text/css" href="../../_static/togglebutton.css" />
    <link rel="stylesheet" type="text/css" href="../../_static/copybutton.css" />
    <link rel="stylesheet" type="text/css" href="../../_static/mystnb.css" />
    <link rel="stylesheet" type="text/css" href="../../_static/sphinx-thebe.css" />
    <link rel="stylesheet" type="text/css" href="../../_static/panels-main.c949a650a448cc0ae9fd3441c0e17fb0.css" />
    <link rel="stylesheet" type="text/css" href="../../_static/panels-variables.06eb56fa6e07937060861dad626602ad.css" />
    
  <link rel="preload" as="script" href="../../_static/js/index.be7d3bbb2ef33a8344ce.js">

    <script data-url_root="../../" id="documentation_options" src="../../_static/documentation_options.js"></script>
    <script src="../../_static/jquery.js"></script>
    <script src="../../_static/underscore.js"></script>
    <script src="../../_static/doctools.js"></script>
    <script src="../../_static/clipboard.min.js"></script>
    <script src="../../_static/copybutton.js"></script>
    <script>let toggleHintShow = 'Click to show';</script>
    <script>let toggleHintHide = 'Click to hide';</script>
    <script>let toggleOpenOnPrint = 'true';</script>
    <script src="../../_static/togglebutton.js"></script>
    <script>var togglebuttonSelector = '.toggle, .admonition.dropdown, .tag_hide_input div.cell_input, .tag_hide-input div.cell_input, .tag_hide_output div.cell_output, .tag_hide-output div.cell_output, .tag_hide_cell.cell, .tag_hide-cell.cell';</script>
    <script src="../../_static/sphinx-book-theme.d59cb220de22ca1c485ebbdc042f0030.js"></script>
    <script>const THEBE_JS_URL = "https://unpkg.com/thebe@0.8.2/lib/index.js"
const thebe_selector = ".thebe,.cell"
const thebe_selector_input = "pre"
const thebe_selector_output = ".output, .cell_output"
</script>
    <script async="async" src="../../_static/sphinx-thebe.js"></script>
    <script>window.MathJax = {"tex": {"macros": {"covMat": "\\boldsymbol{\\Sigma}", "data": "\\mathcal{D}", "defeq": "\\triangleq", "diag": "\\mathrm{diag}", "discreteState": "s", "dotstar": "\\odot", "dynamicsFn": "\\mathbf{f}", "floor": ["\\lfloor#1\\rfloor", 1], "gainMatrix": "\\mathbf{K}", "gainMatrixReverse": "\\mathbf{G}", "gauss": "\\mathcal{N}", "gaussInfo": "\\mathcal{N}_{\\text{info}}", "hidden": "\\mathbf{x}", "hiddenScalar": "x", "hmmInit": "\\boldsymbol{\\pi}", "hmmInitScalar": "\\pi", "hmmObs": "\\mathbf{B}", "hmmObsScalar": "B", "hmmTrans": "\\mathbf{A}", "hmmTransScalar": "A", "infoMat": "\\precMat", "input": "\\mathbf{u}", "inputs": "\\input", "inv": ["{#1}^{-1}", 1], "keyword": ["\\textbf{#1}", 1], "ldsDyn": "\\mathbf{F}", "ldsDynIn": "\\mathbf{B}", "initMean": "\\boldsymbol{\\mean}_0", "initCov": "\\boldsymbol{\\covMat}_0", "ldsObs": "\\mathbf{H}", "ldsObsIn": "\\mathbf{D}", "ldsTrans": "\\ldsDyn", "ldsTransIn": "\\ldsDynIn", "obsCov": "\\mathbf{R}", "obsNoise": "\\boldsymbol{r}", "map": "\\mathrm{map}", "measurementFn": "\\mathbf{h}", "mean": "\\boldsymbol{\\mu}", "mle": "\\mathrm{mle}", "nlatents": "n_x", "nhidden": "\\nlatents", "ninputs": "n_u", "nobs": "n_y", "nsymbols": "n_y", "nstates": "n_s", "obs": "\\mathbf{y}", "obsScalar": "y", "observed": "\\obs", "obsFn": "\\measurementFn", "params": "\\boldsymbol{\\theta}", "precMean": "\\boldsymbol{\\eta}", "precMat": "\\boldsymbol{\\Lambda}", "real": "\\mathbb{R}", "sigmoid": "\\sigma", "softmax": "\\boldsymbol{\\sigma}", "trans": "\\mathsf{T}", "transpose": ["{#1}^{\\trans}", 1], "transCov": "\\mathbf{Q}", "transNoise": "\\mathbf{q}", "valpha": "\\boldsymbol{\\alpha}", "vbeta": "\\boldsymbol{\\beta}", "vdelta": "\\boldsymbol{\\delta}", "vepsilon": "\\boldsymbol{\\epsilon}", "vlambda": "\\boldsymbol{\\lambda}", "vLambda": "\\boldsymbol{\\Lambda}", "vmu": "\\boldsymbol{\\mu}", "vpi": "\\boldsymbol{\\pi}", "vsigma": "\\boldsymbol{\\sigma}", "vSigma": "\\boldsymbol{\\Sigma}", "vone": "\\boldsymbol{1}", "vzero": "\\boldsymbol{0}"}}, "options": {"processHtmlClass": "tex2jax_process|mathjax_process|math|output_area"}}</script>
    <script defer="defer" src="https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js"></script>
    <link rel="index" title="Index" href="../../genindex.html" />
    <link rel="search" title="Search" href="../../search.html" />
    <link rel="next" title="HMM smoothing (forwards-backwards algorithm)" href="hmm_smoother.html" />
    <link rel="prev" title="Hidden Markov Models" href="hmm_index.html" />
    <meta name="viewport" content="width=device-width, initial-scale=1" />
    <meta name="docsearch:language" content="None">
    

    <!-- Google Analytics -->
    
  </head>
  <body data-spy="scroll" data-target="#bd-toc-nav" data-offset="80">
    
    <div class="container-fluid" id="banner"></div>

    
    <div class="container-xl">
      <div class="row">
          
<div class="col-12 col-md-3 bd-sidebar site-navigation show" id="site-navigation">
    
        <div class="navbar-brand-box">
    <a class="navbar-brand text-wrap" href="../../index.html">
      
      
      <h1 class="site-logo" id="site-title">State Space Models: A Modern Approach</h1>
      
    </a>
</div><form class="bd-search d-flex align-items-center" action="../../search.html" method="get">
  <i class="icon fas fa-search"></i>
  <input type="search" class="form-control" name="q" id="search-input" placeholder="Search this book..." aria-label="Search this book..." autocomplete="off" >
</form><nav class="bd-links" id="bd-docs-nav" aria-label="Main">
    <div class="bd-toc-item active">
        <ul class="nav bd-sidenav">
 <li class="toctree-l1">
  <a class="reference internal" href="../../root.html">
   State Space Models: A Modern Approach
  </a>
 </li>
</ul>
<ul class="current nav bd-sidenav">
 <li class="toctree-l1 has-children">
  <a class="reference internal" href="../ssm/ssm_index.html">
   State Space Models
  </a>
  <input class="toctree-checkbox" id="toctree-checkbox-1" name="toctree-checkbox-1" type="checkbox"/>
  <label for="toctree-checkbox-1">
   <i class="fas fa-chevron-down">
   </i>
  </label>
  <ul>
   <li class="toctree-l2">
    <a class="reference internal" href="../ssm/ssm_intro.html">
     What are State Space Models?
    </a>
   </li>
   <li class="toctree-l2">
    <a class="reference internal" href="../ssm/hmm.html">
     Hidden Markov Models
    </a>
   </li>
   <li class="toctree-l2">
    <a class="reference internal" href="../ssm/lds.html">
     Linear Gaussian SSMs
    </a>
   </li>
   <li class="toctree-l2">
    <a class="reference internal" href="../ssm/nlds.html">
     Nonlinear Gaussian SSMs
    </a>
   </li>
   <li class="toctree-l2">
    <a class="reference internal" href="../ssm/inference.html">
     States estimation (inference)
    </a>
   </li>
   <li class="toctree-l2">
    <a class="reference internal" href="../ssm/learning.html">
     Parameter estimation (learning)
    </a>
   </li>
  </ul>
 </li>
 <li class="toctree-l1 current active has-children">
  <a class="reference internal" href="hmm_index.html">
   Hidden Markov Models
  </a>
  <input checked="" class="toctree-checkbox" id="toctree-checkbox-2" name="toctree-checkbox-2" type="checkbox"/>
  <label for="toctree-checkbox-2">
   <i class="fas fa-chevron-down">
   </i>
  </label>
  <ul class="current">
   <li class="toctree-l2 current active">
    <a class="current reference internal" href="#">
     HMM filtering (forwards algorithm)
    </a>
   </li>
   <li class="toctree-l2">
    <a class="reference internal" href="hmm_smoother.html">
     HMM smoothing (forwards-backwards algorithm)
    </a>
   </li>
   <li class="toctree-l2">
    <a class="reference internal" href="hmm_viterbi.html">
     Viterbi algorithm
    </a>
   </li>
   <li class="toctree-l2">
    <a class="reference internal" href="hmm_parallel.html">
     Parallel HMM  smoothing
    </a>
   </li>
   <li class="toctree-l2">
    <a class="reference internal" href="hmm_sampling.html">
     Forwards-filtering backwards-sampling algorithm
    </a>
   </li>
  </ul>
 </li>
 <li class="toctree-l1 has-children">
  <a class="reference internal" href="../lgssm/lgssm_index.html">
   Linear-Gaussian SSMs
  </a>
  <input class="toctree-checkbox" id="toctree-checkbox-3" name="toctree-checkbox-3" type="checkbox"/>
  <label for="toctree-checkbox-3">
   <i class="fas fa-chevron-down">
   </i>
  </label>
  <ul>
   <li class="toctree-l2">
    <a class="reference internal" href="../lgssm/kalman_filter.html">
     Kalman filtering
    </a>
   </li>
   <li class="toctree-l2">
    <a class="reference internal" href="../lgssm/kalman_smoother.html">
     Kalman (RTS) smoother
    </a>
   </li>
   <li class="toctree-l2">
    <a class="reference internal" href="../lgssm/kalman_parallel.html">
     Parallel Kalman Smoother
    </a>
   </li>
   <li class="toctree-l2">
    <a class="reference internal" href="../lgssm/kalman_sampling.html">
     Forwards-filtering backwards sampling
    </a>
   </li>
  </ul>
 </li>
 <li class="toctree-l1 has-children">
  <a class="reference internal" href="../extended/extended_index.html">
   Extended (linearized) methods
  </a>
  <input class="toctree-checkbox" id="toctree-checkbox-4" name="toctree-checkbox-4" type="checkbox"/>
  <label for="toctree-checkbox-4">
   <i class="fas fa-chevron-down">
   </i>
  </label>
  <ul>
   <li class="toctree-l2">
    <a class="reference internal" href="../extended/extended_filter.html">
     Extended Kalman filtering
    </a>
   </li>
   <li class="toctree-l2">
    <a class="reference internal" href="../extended/extended_smoother.html">
     Extended Kalman smoother
    </a>
   </li>
   <li class="toctree-l2">
    <a class="reference internal" href="../extended/extended_parallel.html">
     Parallel extended Kalman smoothing
    </a>
   </li>
  </ul>
 </li>
 <li class="toctree-l1 has-children">
  <a class="reference internal" href="../unscented/unscented_index.html">
   Unscented methods
  </a>
  <input class="toctree-checkbox" id="toctree-checkbox-5" name="toctree-checkbox-5" type="checkbox"/>
  <label for="toctree-checkbox-5">
   <i class="fas fa-chevron-down">
   </i>
  </label>
  <ul>
   <li class="toctree-l2">
    <a class="reference internal" href="../unscented/unscented_filter.html">
     Unscented filtering
    </a>
   </li>
   <li class="toctree-l2">
    <a class="reference internal" href="../unscented/unscented_smoother.html">
     Unscented smoothing
    </a>
   </li>
  </ul>
 </li>
 <li class="toctree-l1">
  <a class="reference internal" href="../quadrature/quadrature_index.html">
   Quadrature and cubature methods
  </a>
 </li>
 <li class="toctree-l1">
  <a class="reference internal" href="../postlin/postlin_index.html">
   Posterior linearization
  </a>
 </li>
 <li class="toctree-l1">
  <a class="reference internal" href="../adf/adf_index.html">
   Assumed Density Filtering
  </a>
 </li>
 <li class="toctree-l1">
  <a class="reference internal" href="../vi/vi_index.html">
   Variational inference
  </a>
 </li>
 <li class="toctree-l1">
  <a class="reference internal" href="../pf/pf_index.html">
   Particle filtering
  </a>
 </li>
 <li class="toctree-l1">
  <a class="reference internal" href="../smc/smc_index.html">
   Sequential Monte Carlo
  </a>
 </li>
 <li class="toctree-l1">
  <a class="reference internal" href="../learning/learning_index.html">
   Offline parameter estimation (learning)
  </a>
 </li>
 <li class="toctree-l1">
  <a class="reference internal" href="../tracking/tracking_index.html">
   Multi-target tracking
  </a>
 </li>
 <li class="toctree-l1">
  <a class="reference internal" href="../ensemble/ensemble_index.html">
   Data assimilation using Ensemble Kalman filter
  </a>
 </li>
 <li class="toctree-l1">
  <a class="reference internal" href="../bnp/bnp_index.html">
   Bayesian non-parametric SSMs
  </a>
 </li>
 <li class="toctree-l1">
  <a class="reference internal" href="../changepoint/changepoint_index.html">
   Changepoint detection
  </a>
 </li>
 <li class="toctree-l1">
  <a class="reference internal" href="../timeseries/timeseries_index.html">
   Timeseries forecasting
  </a>
 </li>
 <li class="toctree-l1">
  <a class="reference internal" href="../gp/gp_index.html">
   Markovian Gaussian processes
  </a>
 </li>
 <li class="toctree-l1">
  <a class="reference internal" href="../ode/ode_index.html">
   Differential equations and SSMs
  </a>
 </li>
 <li class="toctree-l1">
  <a class="reference internal" href="../control/control_index.html">
   Optimal control
  </a>
 </li>
 <li class="toctree-l1">
  <a class="reference internal" href="../../bib.html">
   Bibliography
  </a>
 </li>
</ul>

    </div>
</nav> <!-- To handle the deprecated key -->

<div class="navbar_extra_footer">
  Powered by <a href="https://jupyterbook.org">Jupyter Book</a>
</div>

</div>


<main class="col py-md-3 pl-md-4 bd-content overflow-auto" role="main">
    
    <div class="topbar container-xl fixed-top">
    <div class="topbar-contents row">
        <div class="col-12 col-md-3 bd-topbar-whitespace site-navigation show"></div>
        <div class="col pl-md-4 topbar-main">
            
            <button id="navbar-toggler" class="navbar-toggler ml-0" type="button" data-toggle="collapse"
                data-toggle="tooltip" data-placement="bottom" data-target=".site-navigation" aria-controls="navbar-menu"
                aria-expanded="true" aria-label="Toggle navigation" aria-controls="site-navigation"
                title="Toggle navigation" data-toggle="tooltip" data-placement="left">
                <i class="fas fa-bars"></i>
                <i class="fas fa-arrow-left"></i>
                <i class="fas fa-arrow-up"></i>
            </button>
            
            
<div class="dropdown-buttons-trigger">
    <button id="dropdown-buttons-trigger" class="btn btn-secondary topbarbtn" aria-label="Download this page"><i
            class="fas fa-download"></i></button>

    <div class="dropdown-buttons">
        <!-- ipynb file if we had a myst markdown file -->
        
        <!-- Download raw file -->
        <a class="dropdown-buttons" href="../../_sources/chapters/hmm/hmm_filter.ipynb"><button type="button"
                class="btn btn-secondary topbarbtn" title="Download source file" data-toggle="tooltip"
                data-placement="left">.ipynb</button></a>
        <!-- Download PDF via print -->
        <button type="button" id="download-print" class="btn btn-secondary topbarbtn" title="Print to PDF"
                onclick="printPdf(this)" data-toggle="tooltip" data-placement="left">.pdf</button>
    </div>
</div>

            <!-- Source interaction buttons -->

<div class="dropdown-buttons-trigger">
    <button id="dropdown-buttons-trigger" class="btn btn-secondary topbarbtn"
        aria-label="Connect with source repository"><i class="fab fa-github"></i></button>
    <div class="dropdown-buttons sourcebuttons">
        <a class="repository-button"
            href="https://github.com/probml/ssm-book"><button type="button" class="btn btn-secondary topbarbtn"
                data-toggle="tooltip" data-placement="left" title="Source repository"><i
                    class="fab fa-github"></i>repository</button></a>
        <a class="issues-button"
            href="https://github.com/probml/ssm-book/issues/new?title=Issue%20on%20page%20%2Fchapters/hmm/hmm_filter.html&body=Your%20issue%20content%20here."><button
                type="button" class="btn btn-secondary topbarbtn" data-toggle="tooltip" data-placement="left"
                title="Open an issue"><i class="fas fa-lightbulb"></i>open issue</button></a>
        
    </div>
</div>

            <!-- Full screen (wrap in <a> to have style consistency -->

<a class="full-screen-button"><button type="button" class="btn btn-secondary topbarbtn" data-toggle="tooltip"
        data-placement="bottom" onclick="toggleFullScreen()" aria-label="Fullscreen mode"
        title="Fullscreen mode"><i
            class="fas fa-expand"></i></button></a>

            <!-- Launch buttons -->

<div class="dropdown-buttons-trigger">
    <button id="dropdown-buttons-trigger" class="btn btn-secondary topbarbtn"
        aria-label="Launch interactive content"><i class="fas fa-rocket"></i></button>
    <div class="dropdown-buttons">
        
        <a class="binder-button" href="https://mybinder.org/v2/gh/probml/ssm-book/main?urlpath=tree/chapters/hmm/hmm_filter.ipynb"><button type="button"
                class="btn btn-secondary topbarbtn" title="Launch Binder" data-toggle="tooltip"
                data-placement="left"><img class="binder-button-logo"
                    src="../../_static/images/logo_binder.svg"
                    alt="Interact on binder">Binder</button></a>
        
        
        <a class="colab-button" href="https://colab.research.google.com/github/probml/ssm-book/blob/main/chapters/hmm/hmm_filter.ipynb"><button type="button" class="btn btn-secondary topbarbtn"
                title="Launch Colab" data-toggle="tooltip" data-placement="left"><img class="colab-button-logo"
                    src="../../_static/images/logo_colab.png"
                    alt="Interact on Colab">Colab</button></a>
        
        
    </div>
</div>

        </div>

        <!-- Table of contents -->
        <div class="d-none d-md-block col-md-2 bd-toc show noprint">
            
            <div class="tocsection onthispage pt-5 pb-3">
                <i class="fas fa-list"></i> Contents
            </div>
            <nav id="bd-toc-nav" aria-label="Page">
                <ul class="visible nav section-nav flex-column">
 <li class="toc-h2 nav-item toc-entry">
  <a class="reference internal nav-link" href="#introduction">
   Introduction
  </a>
 </li>
 <li class="toc-h2 nav-item toc-entry">
  <a class="reference internal nav-link" href="#example">
   Example
  </a>
 </li>
 <li class="toc-h2 nav-item toc-entry">
  <a class="reference internal nav-link" href="#normalization-constants">
   Normalization constants
  </a>
 </li>
 <li class="toc-h2 nav-item toc-entry">
  <a class="reference internal nav-link" href="#naive-implementation">
   Naive implementation
  </a>
 </li>
 <li class="toc-h2 nav-item toc-entry">
  <a class="reference internal nav-link" href="#numerically-stable-implementation">
   Numerically stable implementation
  </a>
 </li>
</ul>

            </nav>
        </div>
    </div>
</div>
    <div id="main-content" class="row">
        <div class="col-12 col-md-9 pl-md-3 pr-md-0">
            <!-- Table of contents that is only displayed when printing the page -->
            <div id="jb-print-docs-body" class="onlyprint">
                <h1>HMM filtering (forwards algorithm)</h1>
                <!-- Table of contents -->
                <div id="print-main-content">
                    <div id="jb-print-toc">
                        
                        <div>
                            <h2> Contents </h2>
                        </div>
                        <nav aria-label="Page">
                            <ul class="visible nav section-nav flex-column">
 <li class="toc-h2 nav-item toc-entry">
  <a class="reference internal nav-link" href="#introduction">
   Introduction
  </a>
 </li>
 <li class="toc-h2 nav-item toc-entry">
  <a class="reference internal nav-link" href="#example">
   Example
  </a>
 </li>
 <li class="toc-h2 nav-item toc-entry">
  <a class="reference internal nav-link" href="#normalization-constants">
   Normalization constants
  </a>
 </li>
 <li class="toc-h2 nav-item toc-entry">
  <a class="reference internal nav-link" href="#naive-implementation">
   Naive implementation
  </a>
 </li>
 <li class="toc-h2 nav-item toc-entry">
  <a class="reference internal nav-link" href="#numerically-stable-implementation">
   Numerically stable implementation
  </a>
 </li>
</ul>

                        </nav>
                    </div>
                </div>
            </div>
            
              <div>
                
  <div class="tex2jax_ignore mathjax_ignore section" id="hmm-filtering-forwards-algorithm">
<span id="sec-forwards"></span><h1>HMM filtering (forwards algorithm)<a class="headerlink" href="#hmm-filtering-forwards-algorithm" title="Permalink to this headline">¶</a></h1>
<div class="cell docutils container">
<div class="cell_input docutils container">
<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="c1">### Import standard libraries</span>

<span class="kn">import</span> <span class="nn">abc</span>
<span class="kn">from</span> <span class="nn">dataclasses</span> <span class="kn">import</span> <span class="n">dataclass</span>
<span class="kn">import</span> <span class="nn">functools</span>
<span class="kn">from</span> <span class="nn">functools</span> <span class="kn">import</span> <span class="n">partial</span>
<span class="kn">import</span> <span class="nn">itertools</span>
<span class="kn">import</span> <span class="nn">matplotlib.pyplot</span> <span class="k">as</span> <span class="nn">plt</span>
<span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Any</span><span class="p">,</span> <span class="n">Callable</span><span class="p">,</span> <span class="n">NamedTuple</span><span class="p">,</span> <span class="n">Optional</span><span class="p">,</span> <span class="n">Union</span><span class="p">,</span> <span class="n">Tuple</span>

<span class="kn">import</span> <span class="nn">jax</span>
<span class="kn">import</span> <span class="nn">jax.numpy</span> <span class="k">as</span> <span class="nn">jnp</span>
<span class="kn">from</span> <span class="nn">jax</span> <span class="kn">import</span> <span class="n">lax</span><span class="p">,</span> <span class="n">vmap</span><span class="p">,</span> <span class="n">jit</span><span class="p">,</span> <span class="n">grad</span>
<span class="c1">#from jax.scipy.special import logit</span>
<span class="c1">#from jax.nn import softmax</span>
<span class="kn">import</span> <span class="nn">jax.random</span> <span class="k">as</span> <span class="nn">jr</span>


<span class="kn">import</span> <span class="nn">distrax</span>
<span class="kn">import</span> <span class="nn">optax</span>

<span class="kn">import</span> <span class="nn">jsl</span>
<span class="kn">import</span> <span class="nn">ssm_jax</span>
</pre></div>
</div>
</div>
</div>
<div class="section" id="introduction">
<h2>Introduction<a class="headerlink" href="#introduction" title="Permalink to this headline">¶</a></h2>
<p>The  <span class="math notranslate nohighlight">\(\keyword{Bayes filter}\)</span> is an algorithm for recursively computing
the belief state
<span class="math notranslate nohighlight">\(p(\hidden_t|\obs_{1:t})\)</span> given
the prior belief from the previous step,
<span class="math notranslate nohighlight">\(p(\hidden_{t-1}|\obs_{1:t-1})\)</span>,
the new observation <span class="math notranslate nohighlight">\(\obs_t\)</span>,
and the model.
This can be done using <span class="math notranslate nohighlight">\(\keyword{sequential Bayesian updating}\)</span>.
For a dynamical model, this reduces to the
<span class="math notranslate nohighlight">\(\keyword{predict-update}\)</span> cycle described below.</p>
<p>The <span class="math notranslate nohighlight">\(\keyword{prediction step}\)</span> is just the <span class="math notranslate nohighlight">\(\keyword{Chapman-Kolmogorov equation}\)</span>:</p>
<div class="amsmath math notranslate nohighlight" id="equation-02a2ec0d-8c24-4270-a3c5-b68a6d04b9be">
<span class="eqno">(21)<a class="headerlink" href="#equation-02a2ec0d-8c24-4270-a3c5-b68a6d04b9be" title="Permalink to this equation">¶</a></span>\[\begin{align}
p(\hidden_t|\obs_{1:t-1})
= \int p(\hidden_t|\hidden_{t-1}) p(\hidden_{t-1}|\obs_{1:t-1}) d\hidden_{t-1}
\end{align}\]</div>
<p>The prediction step computes
the <span class="math notranslate nohighlight">\(\keyword{one-step-ahead predictive distribution}\)</span>
for the latent state, which converts
the posterior from the previous time step to become the prior
for the current step.</p>
<p>The <span class="math notranslate nohighlight">\(\keyword{update step}\)</span>
is just Bayes rule:</p>
<div class="amsmath math notranslate nohighlight" id="equation-be8910aa-7596-460f-9268-280f22fdce5e">
<span class="eqno">(22)<a class="headerlink" href="#equation-be8910aa-7596-460f-9268-280f22fdce5e" title="Permalink to this equation">¶</a></span>\[\begin{align}
p(\hidden_t|\obs_{1:t}) = \frac{1}{Z_t}
p(\obs_t|\hidden_t) p(\hidden_t|\obs_{1:t-1})
\end{align}\]</div>
<p>where the normalization constant is</p>
<div class="amsmath math notranslate nohighlight" id="equation-18e83263-589e-4859-84ea-9721705c1f4e">
<span class="eqno">(23)<a class="headerlink" href="#equation-18e83263-589e-4859-84ea-9721705c1f4e" title="Permalink to this equation">¶</a></span>\[\begin{align}
Z_t = \int p(\obs_t|\hidden_t) p(\hidden_t|\obs_{1:t-1}) d\hidden_{t}
= p(\obs_t|\obs_{1:t-1})
\end{align}\]</div>
<p>Note that we can derive the log marginal likelihood from these normalization constants
as follows:</p>
<div class="math notranslate nohighlight" id="equation-eqn-logz">
<span class="eqno">(24)<a class="headerlink" href="#equation-eqn-logz" title="Permalink to this equation">¶</a></span>\[\log p(\obs_{1:T})
= \sum_{t=1}^{T} \log p(\obs_t|\obs_{1:t-1})
= \sum_{t=1}^{T} \log Z_t\]</div>
<p>When the latent states <span class="math notranslate nohighlight">\(\hidden_t\)</span> are discrete, as in HMM,
the above integrals become sums.
In particular, suppose we define
the <span class="math notranslate nohighlight">\(\keyword{belief state}\)</span> as <span class="math notranslate nohighlight">\(\alpha_t(j) \defeq p(\hidden_t=j|\obs_{1:t})\)</span>,
the  <span class="math notranslate nohighlight">\(\keyword{local evidence}\)</span> (or <span class="math notranslate nohighlight">\(\keyword{local likelihood}\)</span>)
as <span class="math notranslate nohighlight">\(\lambda_t(j) \defeq p(\obs_t|\hidden_t=j)\)</span>,
and the transition matrix as
<span class="math notranslate nohighlight">\(\hmmTrans(i,j)  = p(\hidden_t=j|\hidden_{t-1}=i)\)</span>.
Then the predict step becomes</p>
<div class="math notranslate nohighlight" id="equation-eqn-predictivehmm">
<span class="eqno">(25)<a class="headerlink" href="#equation-eqn-predictivehmm" title="Permalink to this equation">¶</a></span>\[\alpha_{t|t-1}(j) \defeq p(\hidden_t=j|\obs_{1:t-1})
 = \sum_i \alpha_{t-1}(i) A(i,j)\]</div>
<p>and the update step becomes</p>
<div class="math notranslate nohighlight" id="equation-eqn-fwdseqn">
<span class="eqno">(26)<a class="headerlink" href="#equation-eqn-fwdseqn" title="Permalink to this equation">¶</a></span>\[\alpha_t(j)
= \frac{1}{Z_t} \lambda_t(j) \alpha_{t|t-1}(j)
= \frac{1}{Z_t} \lambda_t(j) \left[\sum_i \alpha_{t-1}(i) \hmmTrans(i,j)  \right]\]</div>
<p>where
the  normalization constant for each time step is given by</p>
<div class="math notranslate nohighlight" id="equation-eqn-hmmz">
<span class="eqno">(27)<a class="headerlink" href="#equation-eqn-hmmz" title="Permalink to this equation">¶</a></span>\[\begin{split}\begin{align}
Z_t \defeq p(\obs_t|\obs_{1:t-1})
&amp;=  \sum_{j=1}^K p(\obs_t|\hidden_t=j)  p(\hidden_t=j|\obs_{1:t-1}) \\
&amp;=  \sum_{j=1}^K \lambda_t(j) \alpha_{t|t-1}(j)
\end{align}\end{split}\]</div>
<p>Since all the quantities are finite length vectors and matrices,
we can implement the whole procedure using matrix vector multoplication:</p>
<div class="math notranslate nohighlight" id="equation-eqn-fwdsalgomatrixform">
<span class="eqno">(28)<a class="headerlink" href="#equation-eqn-fwdsalgomatrixform" title="Permalink to this equation">¶</a></span>\[\valpha_t =\text{normalize}\left(
\vlambda_t \dotstar  (\hmmTrans^{\trans} \valpha_{t-1}) \right)\]</div>
<p>where <span class="math notranslate nohighlight">\(\dotstar\)</span> represents
elementwise vector multiplication,
and the <span class="math notranslate nohighlight">\(\text{normalize}\)</span> function just ensures its argument sums to one.</p>
</div>
<div class="section" id="example">
<h2>Example<a class="headerlink" href="#example" title="Permalink to this headline">¶</a></h2>
<p>In <a class="reference internal" href="../ssm/inference.html#sec-casino-inference"><span class="std std-ref">Example: inference in the casino HMM</span></a>
we illustrate
filtering for the casino HMM,
applied to a random sequence <span class="math notranslate nohighlight">\(\obs_{1:T}\)</span> of length <span class="math notranslate nohighlight">\(T=300\)</span>.
In blue, we plot the probability that the dice is in the loaded (vs fair) state,
based on the evidence seen so far.
The gray bars indicate time intervals during which the generative
process actually switched to the loaded dice.
We see that the probability generally increases in the right places.</p>
</div>
<div class="section" id="normalization-constants">
<h2>Normalization constants<a class="headerlink" href="#normalization-constants" title="Permalink to this headline">¶</a></h2>
<p>In most publications on HMMs,
such as <span id="id1">[<a class="reference internal" href="../../bib.html#id32" title="L. R. Rabiner. A tutorial on Hidden Markov Models and selected applications in speech recognition. Proc. of the IEEE, 77(2):257–286, 1989.">Rab89</a>]</span>,
the forwards message is defined
as the following unnormalized joint probability:</p>
<div class="math notranslate nohighlight">
\[\alpha'_t(j) = p(\hidden_t=j,\obs_{1:t}) 
= \lambda_t(j) \left[\sum_i \alpha'_{t-1}(i) A(i,j)  \right]\]</div>
<p>In this book we define the forwards message   as the normalized
conditional probability</p>
<div class="math notranslate nohighlight">
\[\alpha_t(j) = p(\hidden_t=j|\obs_{1:t}) 
= \frac{1}{Z_t} \lambda_t(j) \left[\sum_i \alpha_{t-1}(i) A(i,j)  \right]\]</div>
<p>where <span class="math notranslate nohighlight">\(Z_t = p(\obs_t|\obs_{1:t-1})\)</span>.</p>
<p>The “traditional” unnormalized form has several problems.
First, it rapidly suffers from numerical underflow,
since the probability of
the joint event that <span class="math notranslate nohighlight">\((\hidden_t=j,\obs_{1:t})\)</span>
is vanishingly small.
To see why, suppose the observations are independent of the states.
In this case, the unnormalized joint has the form</p>
<div class="amsmath math notranslate nohighlight" id="equation-2f9dfff7-1a23-4760-98f4-68f254fa0022">
<span class="eqno">(29)<a class="headerlink" href="#equation-2f9dfff7-1a23-4760-98f4-68f254fa0022" title="Permalink to this equation">¶</a></span>\[\begin{align}
p(\hidden_t=j,\obs_{1:t}) = p(\hidden_t=j)\prod_{i=1}^t p(\obs_i)
\end{align}\]</div>
<p>which becomes exponentially small with <span class="math notranslate nohighlight">\(t\)</span>, because we multiply
many probabilities which are less than one.
Second, the unnormalized probability is less interpretable,
since it is a joint distribution over states and observations,
rather than a conditional probability of states given observations.
Third, the unnormalized joint form is harder to approximate
than the normalized form.
Of course,
the two definitions only differ by a
multiplicative constant
<span id="id2">[<a class="reference internal" href="../../bib.html#id35" title="Pierre A Devijver. Baum's forward-backward algorithm revisited. Pattern Recognition Letters, 3(6):369–373, 1985.">Dev85</a>]</span>,
so the algorithmic difference is just
one line of code (namely the presence or absence of a call to the <code class="docutils literal notranslate"><span class="pre">normalize</span></code> function).</p>
</div>
<div class="section" id="naive-implementation">
<h2>Naive implementation<a class="headerlink" href="#naive-implementation" title="Permalink to this headline">¶</a></h2>
<p>Below we give a simple numpy implementation of the forwards algorithm.
We assume the HMM uses categorical observations, for simplicity.</p>
<div class="cell docutils container">
<div class="cell_input docutils container">
<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="k">def</span> <span class="nf">normalize_np</span><span class="p">(</span><span class="n">u</span><span class="p">,</span> <span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span> <span class="n">eps</span><span class="o">=</span><span class="mf">1e-15</span><span class="p">):</span>
    <span class="n">u</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">where</span><span class="p">(</span><span class="n">u</span> <span class="o">==</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="n">np</span><span class="o">.</span><span class="n">where</span><span class="p">(</span><span class="n">u</span> <span class="o">&lt;</span> <span class="n">eps</span><span class="p">,</span> <span class="n">eps</span><span class="p">,</span> <span class="n">u</span><span class="p">))</span>
    <span class="n">c</span> <span class="o">=</span> <span class="n">u</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="n">axis</span><span class="p">)</span>
    <span class="n">c</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">where</span><span class="p">(</span><span class="n">c</span> <span class="o">==</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="n">c</span><span class="p">)</span>
    <span class="k">return</span> <span class="n">u</span> <span class="o">/</span> <span class="n">c</span><span class="p">,</span> <span class="n">c</span>

<span class="k">def</span> <span class="nf">hmm_forwards_np</span><span class="p">(</span><span class="n">trans_mat</span><span class="p">,</span> <span class="n">obs_mat</span><span class="p">,</span> <span class="n">init_dist</span><span class="p">,</span> <span class="n">obs_seq</span><span class="p">):</span>
    <span class="n">n_states</span><span class="p">,</span> <span class="n">n_obs</span> <span class="o">=</span> <span class="n">obs_mat</span><span class="o">.</span><span class="n">shape</span>
    <span class="n">seq_len</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">obs_seq</span><span class="p">)</span>

    <span class="n">alpha_hist</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">seq_len</span><span class="p">,</span> <span class="n">n_states</span><span class="p">))</span>
    <span class="n">ll_hist</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">(</span><span class="n">seq_len</span><span class="p">)</span>  <span class="c1"># loglikelihood history</span>

    <span class="n">alpha_n</span> <span class="o">=</span> <span class="n">init_dist</span> <span class="o">*</span> <span class="n">obs_mat</span><span class="p">[:,</span> <span class="n">obs_seq</span><span class="p">[</span><span class="mi">0</span><span class="p">]]</span>
    <span class="n">alpha_n</span><span class="p">,</span> <span class="n">cn</span> <span class="o">=</span> <span class="n">normalize_np</span><span class="p">(</span><span class="n">alpha_n</span><span class="p">)</span>

    <span class="n">alpha_hist</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="n">alpha_n</span>
    <span class="n">log_normalizer</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">log</span><span class="p">(</span><span class="n">cn</span><span class="p">)</span>

    <span class="k">for</span> <span class="n">t</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="n">seq_len</span><span class="p">):</span>
        <span class="n">alpha_n</span> <span class="o">=</span> <span class="n">obs_mat</span><span class="p">[:,</span> <span class="n">obs_seq</span><span class="p">[</span><span class="n">t</span><span class="p">]]</span> <span class="o">*</span> <span class="p">(</span><span class="n">alpha_n</span><span class="p">[:,</span> <span class="kc">None</span><span class="p">]</span> <span class="o">*</span> <span class="n">trans_mat</span><span class="p">)</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span>
        <span class="n">alpha_n</span><span class="p">,</span> <span class="n">zn</span> <span class="o">=</span> <span class="n">normalize_np</span><span class="p">(</span><span class="n">alpha_n</span><span class="p">)</span>

        <span class="n">alpha_hist</span><span class="p">[</span><span class="n">t</span><span class="p">]</span> <span class="o">=</span> <span class="n">alpha_n</span>
        <span class="n">log_normalizer</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">log</span><span class="p">(</span><span class="n">zn</span><span class="p">)</span> <span class="o">+</span> <span class="n">log_normalizer</span>

    <span class="k">return</span>  <span class="n">log_normalizer</span><span class="p">,</span> <span class="n">alpha_hist</span>
</pre></div>
</div>
</div>
</div>
</div>
<div class="section" id="numerically-stable-implementation">
<h2>Numerically stable implementation<a class="headerlink" href="#numerically-stable-implementation" title="Permalink to this headline">¶</a></h2>
<p>In practice it is more numerically stable to compute
the log likelihoods <span class="math notranslate nohighlight">\(\ell_t(j) = \log p(\obs_t|\hidden_t=j)\)</span>,
rather than the likelioods <span class="math notranslate nohighlight">\(\lambda_t(j) = p(\obs_t|\hidden_t=j)\)</span>.
In this case, we can perform the posterior updating in a numerically stable way as follows.
Define <span class="math notranslate nohighlight">\(L_t = \max_j \ell_t(j)\)</span> and</p>
<div class="amsmath math notranslate nohighlight" id="equation-1cf0a390-376b-4da3-873f-daf7e1489cbb">
<span class="eqno">(30)<a class="headerlink" href="#equation-1cf0a390-376b-4da3-873f-daf7e1489cbb" title="Permalink to this equation">¶</a></span>\[\begin{align}
\tilde{p}(\hidden_t=j,\obs_t|\obs_{1:t-1})
&amp;\defeq p(\hidden_t=j|\obs_{1:t-1}) p(\obs_t|\hidden_t=j) e^{-L_t} \\
 &amp;= p(\hidden_t=j|\obs_{1:t-1}) e^{\ell_t(j) - L_t}
\end{align}\]</div>
<p>Then we have</p>
<div class="amsmath math notranslate nohighlight" id="equation-562d5a1d-b724-4093-85d2-ec62f15c5421">
<span class="eqno">(31)<a class="headerlink" href="#equation-562d5a1d-b724-4093-85d2-ec62f15c5421" title="Permalink to this equation">¶</a></span>\[\begin{align}
p(\hidden_t=j|\obs_t,\obs_{1:t-1})
  &amp;= \frac{1}{\tilde{Z}_t} \tilde{p}(\hidden_t=j,\obs_t|\obs_{1:t-1}) \\
\tilde{Z}_t &amp;= \sum_j \tilde{p}(\hidden_t=j,\obs_t|\obs_{1:t-1})
= p(\obs_t|\obs_{1:t-1}) e^{-L_t} \\
\log Z_t &amp;= \log p(\obs_t|\obs_{1:t-1}) = \log \tilde{Z}_t + L_t
\end{align}\]</div>
<p>Below we show some JAX code that implements this core operation.</p>
<div class="cell docutils container">
<div class="cell_input docutils container">
<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="k">def</span> <span class="nf">_condition_on</span><span class="p">(</span><span class="n">probs</span><span class="p">,</span> <span class="n">ll</span><span class="p">):</span>
    <span class="n">ll_max</span> <span class="o">=</span> <span class="n">ll</span><span class="o">.</span><span class="n">max</span><span class="p">()</span>
    <span class="n">new_probs</span> <span class="o">=</span> <span class="n">probs</span> <span class="o">*</span> <span class="n">jnp</span><span class="o">.</span><span class="n">exp</span><span class="p">(</span><span class="n">ll</span> <span class="o">-</span> <span class="n">ll_max</span><span class="p">)</span>
    <span class="n">norm</span> <span class="o">=</span> <span class="n">new_probs</span><span class="o">.</span><span class="n">sum</span><span class="p">()</span>
    <span class="n">new_probs</span> <span class="o">/=</span> <span class="n">norm</span>
    <span class="n">log_norm</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">log</span><span class="p">(</span><span class="n">norm</span><span class="p">)</span> <span class="o">+</span> <span class="n">ll_max</span>
    <span class="k">return</span> <span class="n">new_probs</span><span class="p">,</span> <span class="n">log_norm</span>
</pre></div>
</div>
</div>
</div>
<p>With the above function, we can implement a more numerically stable version of the forwards filter,
that works for any likelihood function, as shown below. It takes in the prior predictive distribution,
<span class="math notranslate nohighlight">\(\alpha_{t|t-1}\)</span>,
stored in <code class="docutils literal notranslate"><span class="pre">predicted_probs</span></code>, and conditions them on the log-likelihood for each time step <span class="math notranslate nohighlight">\(\ell_t\)</span> to get the
posterior, <span class="math notranslate nohighlight">\(\alpha_t\)</span>, stored in <code class="docutils literal notranslate"><span class="pre">filtered_probs</span></code>,
which is then converted to the prediction for the next state, <span class="math notranslate nohighlight">\(\alpha_{t+1|t}\)</span>.</p>
<div class="cell docutils container">
<div class="cell_input docutils container">
<div class="highlight-ipython3 notranslate"><div class="highlight"><pre><span></span><span class="k">def</span> <span class="nf">_predict</span><span class="p">(</span><span class="n">probs</span><span class="p">,</span> <span class="n">A</span><span class="p">):</span>
    <span class="k">return</span> <span class="n">A</span><span class="o">.</span><span class="n">T</span> <span class="o">@</span> <span class="n">probs</span>


<span class="k">def</span> <span class="nf">hmm_filter</span><span class="p">(</span><span class="n">initial_distribution</span><span class="p">,</span>
               <span class="n">transition_matrix</span><span class="p">,</span>
               <span class="n">log_likelihoods</span><span class="p">):</span>
    <span class="k">def</span> <span class="nf">_step</span><span class="p">(</span><span class="n">carry</span><span class="p">,</span> <span class="n">t</span><span class="p">):</span>
        <span class="n">log_normalizer</span><span class="p">,</span> <span class="n">predicted_probs</span> <span class="o">=</span> <span class="n">carry</span>

        <span class="c1"># Get parameters for time t</span>
        <span class="n">get</span> <span class="o">=</span> <span class="k">lambda</span> <span class="n">x</span><span class="p">:</span> <span class="n">x</span><span class="p">[</span><span class="n">t</span><span class="p">]</span> <span class="k">if</span> <span class="n">x</span><span class="o">.</span><span class="n">ndim</span> <span class="o">==</span> <span class="mi">3</span> <span class="k">else</span> <span class="n">x</span>
        <span class="n">A</span> <span class="o">=</span> <span class="n">get</span><span class="p">(</span><span class="n">transition_matrix</span><span class="p">)</span>
        <span class="n">ll</span> <span class="o">=</span> <span class="n">log_likelihoods</span><span class="p">[</span><span class="n">t</span><span class="p">]</span>

        <span class="c1"># Condition on emissions at time t, being careful not to overflow</span>
        <span class="n">filtered_probs</span><span class="p">,</span> <span class="n">log_norm</span> <span class="o">=</span> <span class="n">_condition_on</span><span class="p">(</span><span class="n">predicted_probs</span><span class="p">,</span> <span class="n">ll</span><span class="p">)</span>
        <span class="c1"># Update the log normalizer</span>
        <span class="n">log_normalizer</span> <span class="o">+=</span> <span class="n">log_norm</span>
        <span class="c1"># Predict the next state</span>
        <span class="n">predicted_probs</span> <span class="o">=</span> <span class="n">_predict</span><span class="p">(</span><span class="n">filtered_probs</span><span class="p">,</span> <span class="n">A</span><span class="p">)</span>

        <span class="k">return</span> <span class="p">(</span><span class="n">log_normalizer</span><span class="p">,</span> <span class="n">predicted_probs</span><span class="p">),</span> <span class="p">(</span><span class="n">filtered_probs</span><span class="p">,</span> <span class="n">predicted_probs</span><span class="p">)</span>

    <span class="n">num_timesteps</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="n">log_likelihoods</span><span class="p">)</span>
    <span class="n">carry</span> <span class="o">=</span> <span class="p">(</span><span class="mf">0.0</span><span class="p">,</span> <span class="n">initial_distribution</span><span class="p">)</span>
    <span class="p">(</span><span class="n">log_normalizer</span><span class="p">,</span> <span class="n">_</span><span class="p">),</span> <span class="p">(</span><span class="n">filtered_probs</span><span class="p">,</span> <span class="n">predicted_probs</span><span class="p">)</span> <span class="o">=</span> <span class="n">lax</span><span class="o">.</span><span class="n">scan</span><span class="p">(</span>
        <span class="n">_step</span><span class="p">,</span> <span class="n">carry</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="n">num_timesteps</span><span class="p">))</span>
    <span class="k">return</span> <span class="n">log_normalizer</span><span class="p">,</span> <span class="n">filtered_probs</span><span class="p">,</span> <span class="n">predicted_probs</span>
</pre></div>
</div>
</div>
</div>
<p>TODO: check equivalence of these two implementations!</p>
</div>
</div>

    <script type="text/x-thebe-config">
    {
        requestKernel: true,
        binderOptions: {
            repo: "binder-examples/jupyter-stacks-datascience",
            ref: "master",
        },
        codeMirrorConfig: {
            theme: "abcdef",
            mode: "python"
        },
        kernelOptions: {
            kernelName: "python3",
            path: "./chapters/hmm"
        },
        predefinedOutput: true
    }
    </script>
    <script>kernelName = 'python3'</script>

              </div>
              
            
                <!-- Previous / next buttons -->
<div class='prev-next-area'> 
    <a class='left-prev' id="prev-link" href="hmm_index.html" title="previous page">
        <i class="fas fa-angle-left"></i>
        <div class="prev-next-info">
            <p class="prev-next-subtitle">previous</p>
            <p class="prev-next-title">Hidden Markov Models</p>
        </div>
    </a>
    <a class='right-next' id="next-link" href="hmm_smoother.html" title="next page">
    <div class="prev-next-info">
        <p class="prev-next-subtitle">next</p>
        <p class="prev-next-title">HMM smoothing (forwards-backwards algorithm)</p>
    </div>
    <i class="fas fa-angle-right"></i>
    </a>
</div>
            
        </div>
    </div>
    <footer class="footer">
  <p>
    
      By Kevin Murphy, Scott Linderman, et al.<br/>
    
        &copy; Copyright 2021.<br/>
  </p>
</footer>
</main>


      </div>
    </div>
  
  <script src="../../_static/js/index.be7d3bbb2ef33a8344ce.js"></script>

  </body>
</html>