### Install Dynamax

Source: https://github.com/probml/dynamax/blob/main/dynamax/hidden_markov_model/demos/multinomial_hmm.ipynb

Installs the Dynamax library from its GitHub repository if it's not already installed. This is a prerequisite for running the examples.

```python
try:
    import dynamax
except ModuleNotFoundError:
    %pip install -qq git+https://github.com/probml/dynamax.git
    import dynamax
```

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

### Install dynamax

Source: https://github.com/probml/dynamax/blob/main/docs/notebooks/linear_gaussian_ssm/kf_linreg.ipynb

Installs the dynamax library with notebook support if it's not already installed. This is a prerequisite for running the subsequent code examples.

```python
%%capture
try:
    import dynamax
except ModuleNotFoundError:
    print('installing dynamax')
    %pip install dynamax[notebooks]
    import dynamax
```

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

### Install Dynamax in Development Mode

Source: https://github.com/probml/dynamax/blob/main/CONTRIBUTING.md

Clone the repository and install Dynamax with development dependencies. This setup is necessary for making and testing code changes.

```bash
git clone https://github.com/your-username-here/dynamax.git
cd dynamax
pip install -e '.[dev]'
```

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

### Install dynamax with notebook support

Source: https://github.com/probml/dynamax/blob/main/docs/notebooks/generalized_gaussian_ssm/cmgf_poisson_demo.ipynb

Installs the dynamax library with necessary dependencies for notebook usage. It checks if dynamax is already installed and installs it if not.

```python
%%capture
try:
    import dynamax
except ModuleNotFoundError:
    print('installing dynamax')
    %pip install -q dynamax[notebooks]
    import dynamax
```

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

### Install Dynamax

Source: https://github.com/probml/dynamax/blob/main/dynamax/hidden_markov_model/demos/fixed_lag_smoother_hmm.ipynb

Installs the Dynamax library if it's not already found. This is a prerequisite for using the library's functionalities.

```python
try:
    import dynamax
except ModuleNotFoundError:
    %pip install git+https://github.com/probml/dynamax.git
    import dynamax
```

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

### Install Dynamax from PyPI

Source: https://github.com/probml/dynamax/blob/main/docs/index.rst

Installs the latest release of dynamax and core dependencies. Use `[notebooks]` for demo notebook dependencies.

```console
pip install dynamax                 # Install dynamax and core dependencies, or
pip install dynamax[notebooks]      # Install with demo notebook dependencies
```

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

### Install Dynamax

Source: https://github.com/probml/dynamax/blob/main/dynamax/hidden_markov_model/demos/low_rank_gaussian_hmm.ipynb

Installs the dynamax library if it is not already present. This is a prerequisite for using the HMM functionalities.

```python
try:
    import dynamax
except ModuleNotFoundError:
    print('installing dynamax')
    %pip install -qq git+https://github.com/probml/dynamax.git
    import dynamax
```

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

### MLP Instantiation and Setup

Source: https://github.com/probml/dynamax/blob/main/docs/notebooks/nonlinear_gaussian_ssm/ekf_mlp.ipynb

Instantiates an MLP model with specified dimensions and retrieves its flattened initial parameters and apply function. This is a common setup step before training or evaluating the MLP.

```python
input_dim, hidden_dim, output_dim = 1, 6, 1
model_dims = [input_dim, hidden_dim, output_dim]
_, flat_params, _, apply_fn = get_mlp_flattened_params(model_dims)

```

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

### Install Dynamax with Development Dependencies

Source: https://github.com/probml/dynamax/blob/main/docs/index.rst

Installs dynamax along with test and documentation dependencies for developers. Requires cloning the repository first.

```console
git clone git@github.com:probml/dynamax.git
cd dynamax
pip install -e '.[dev]'
```

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

### Install Dynamax Development Branch

Source: https://github.com/probml/dynamax/blob/main/docs/index.rst

Installs the latest development version of dynamax directly from its GitHub repository.

```console
pip install git+https://github.com/probml/dynamax.git
```

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

### Install Latest Dynamax Release

Source: https://github.com/probml/dynamax/blob/main/README.md

Install the latest release of dynamax from PyPI. Use the `[notebooks]` extra to include dependencies for demo notebooks.

```console
pip install dynamax

```

```console
pip install dynamax[notebooks]

```

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

### Install Latest Dynamax Development Branch

Source: https://github.com/probml/dynamax/blob/main/README.md

Install the latest development branch of dynamax directly from its GitHub repository.

```console
pip install git+https://github.com/probml/dynamax.git

```

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

### Install ProbML Utils

Source: https://github.com/probml/dynamax/blob/main/dynamax/hidden_markov_model/demos/bernoulli_hmm_example.ipynb

Installs the probml_utils library if it is not already found. This library provides utility functions for plotting and LaTeX formatting.

```python
try:
    from probml_utils import savefig, latexify, is_latexify_enabled
except ModuleNotFoundError:
    %pip install git+https://github.com/probml/probml-utils.git
    from probml_utils import savefig, latexify, is_latexify_enabled
```

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

### Install Blackjax

Source: https://github.com/probml/dynamax/blob/main/docs/notebooks/linear_gaussian_ssm/lgssm_hmc.ipynb

Installs the blackjax library if not already present. Blackjax is used for MCMC sampling, including HMC, which is essential for Bayesian parameter estimation.

```python
try:
    import blackjax
except ModuleNotFoundError:
    print('installing blackjax')
    %pip install -qq blackjax
    import blackjax
```

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

### Run Tests with Pytest

Source: https://github.com/probml/dynamax/blob/main/CONTRIBUTING.md

Install pytest and run all project tests to ensure your changes do not break existing functionality. This is a crucial step before submitting a pull request.

```bash
pip install pytest
pytest
```

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

### EKF Parameter Setup for Online Inference

Source: https://github.com/probml/dynamax/blob/main/docs/notebooks/nonlinear_gaussian_ssm/ekf_mlp.ipynb

Initializes EKF parameters for online learning. The dynamics function is the identity, and the emission function is the MLP's apply function. Use this to set up the EKF for training an MLP with sequential data.

```python
state_dim, emission_dim = flat_params.size, output_dim
ekf_params = ParamsNLGSSM(
    initial_mean=flat_params,
    initial_covariance=jnp.eye(state_dim) * 100,
    dynamics_function=lambda x, u: x,
    dynamics_covariance=jnp.eye(state_dim) * 1e-4,
    emission_function=apply_fn,
    emission_covariance=jnp.eye(emission_dim) * y_std**2,
)

# Run EKF on training set to train MLP
ekf_post = extended_kalman_filter(ekf_params, emissions, inputs=inputs)
w_means, w_covs = ekf_post.filtered_means, ekf_post.filtered_covariances
```

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

### Import necessary libraries

Source: https://github.com/probml/dynamax/blob/main/dynamax/hidden_markov_model/demos/poisson_hmm_neurons.ipynb

Imports essential libraries for the Poisson HMM example, including JAX for numerical operations, Matplotlib for plotting, and specific modules from dynamax.

```python
from functools import partial
import jax.numpy as jnp
import jax.random as jr
from jax import vmap
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
from scipy.optimize import linear_sum_assignment
import seaborn as sns

from dynamax.hidden_markov_model import PoissonHMM
```

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

### Build Sphinx Documentation

Source: https://github.com/probml/dynamax/blob/main/docs/README.md

Navigate to the docs directory and run 'make html' to build the documentation. This process executes all Jupyter notebooks and integrates their output.

```bash
cd docs/
make html
```

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

### Instantiate and Sample from Poisson GLM-HMM

Source: https://github.com/probml/dynamax/blob/main/docs/notebooks/hmm/custom_hmm.ipynb

Demonstrates how to instantiate a Poisson GLM-HMM, define a custom transition matrix, initialize its parameters, generate inputs, and sample states and emissions from the model.

```python
# Make a Poisson GLM-HMM
init_key, sample_key = jr.split(jr.PRNGKey(0))
num_states = 3
emission_dim = 10
input_dim = 5

# Construct a simple, sticky transition matrix
transition_probs = (jnp.arange(num_states)**5).astype(float)
transition_probs /= transition_probs.sum()
transition_matrix = jnp.zeros((num_states, num_states))
for k, p in enumerate(transition_probs[::-1]):
    transition_matrix += jnp.roll(p * jnp.eye(num_states), k, axis=1)
    
# Construct the Poisson GLM-HMM and 
true_model = PoissonGLMHMM(num_states, emission_dim, input_dim)
true_params, _ = true_model.initialize(key=init_key, transition_matrix=transition_matrix)

# Generate random inputs
num_timesteps = 1000
inputs = tfd.Exponential(rate=1.0).sample(
    seed=sample_key, sample_shape=(num_timesteps, input_dim))

# Sample states and emissions from the model
true_states, emissions = true_model.sample(
    true_params, sample_key, num_timesteps, inputs=inputs)
```

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

### Initialize and Sample LGSSM Parameters

Source: https://github.com/probml/dynamax/blob/main/docs/notebooks/linear_gaussian_ssm/lgssm_hmc.ipynb

Constructs a true LGSSM model with randomly initialized parameters and samples states and emissions. Requires JAX and dynamax libraries.

```python
state_dim = 2
emission_dim = 10
num_timesteps = 100
k1, k2, k3 = jr.split(jr.PRNGKey(0), 3)

# Construct the true model with randomly initialized parameters
true_A = 0.99 * random_rotation(seed=k1, n=state_dim, theta=jnp.pi / 10)
true_Sigma = 0.01 * jnp.eye(state_dim)
true_model = LinearGaussianConjugateSSM(state_dim, emission_dim)
true_params, param_props = true_model.initialize(
    key=k1, dynamics_weights=true_A, dynamics_covariance=true_Sigma)

# Sample states and emissions from the true model
true_states, emissions = true_model.sample(true_params, k3, num_timesteps)
```

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

### Simulate and Fit Gaussian HMM with EM

Source: https://github.com/probml/dynamax/blob/main/README.md

Demonstrates initializing a Gaussian HMM, sampling data, fitting the model using EM, and performing posterior inference. Requires JAX and Matplotlib.

```python
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from dynamax.hidden_markov_model import GaussianHMM

key1, key2, key3 = jr.split(jr.PRNGKey(0), 3)
num_states = 3
emission_dim = 2
num_timesteps = 1000

# Make a Gaussian HMM and sample data from it
hmm = GaussianHMM(num_states, emission_dim)
true_params, _ = hmm.initialize(key1)
true_states, emissions = hmm.sample(true_params, key2, num_timesteps)

# Make a new Gaussian HMM and fit it with EM
params, props = hmm.initialize(key3, method="kmeans", emissions=emissions)
params, lls = hmm.fit_em(params, props, emissions, num_iters=20)

# Plot the marginal log probs across EM iterations
plt.plot(lls)
plt.xlabel("EM iterations")
plt.ylabel("marginal log prob.")

# Use fitted model for posterior inference
post = hmm.smoother(params, emissions)
print(post.smoothed_probs.shape) # (1000, 3)
```

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

### Import Flax for MLP

Source: https://github.com/probml/dynamax/blob/main/docs/notebooks/generalized_gaussian_ssm/cmgf_mlp_classification_demo.ipynb

Imports Flax, a neural network library, for defining the MLP classifier. Installs Flax if not found.

```python
try:
    import flax.linen as nn
except ModuleNotFoundError:
    print('installing flax')
    %pip install -qq flax
import flax.linen as nn
```

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

### Initializing and Plotting Random Model Forecasts

Source: https://github.com/probml/dynamax/blob/main/docs/notebooks/linear_gaussian_ssm/lgssm_learning.ipynb

Initializes a new LinearGaussianSSM with random parameters and plots its emissions and forecasts. This demonstrates the performance of an untrained model.

```python
# Plot predictions from a random, untrained model
init_key = jr.PRNGKey(42)
model = LinearGaussianSSM(state_dim, emission_dim)
params, param_props = model.initialize(init_key)
```

```python
plot_emissions_and_forecast(model, params, emissions)
plt.suptitle("randomly initialized model", y=1.02)
```

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

### Importing Dynamax Inference Module

Source: https://github.com/probml/dynamax/blob/main/dynamax/generalized_gaussian_ssm/demos/cmgf_multiclass_logreg_demo.ipynb

Attempts to import inference functions from Dynamax. If the module is not found, it installs Dynamax and then imports the module.

```python
try:
    from dynamax.generalized_gaussian_ssm.inference import *
except ModuleNotFoundError:
    print('installing dynamax')
    %pip install -qq dynamax
    from dynamax.generalized_gaussian_ssm.inference import *
```

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

### Initialize GaussianHMM Model

Source: https://github.com/probml/dynamax/blob/main/docs/index.rst

Initializes a Gaussian Hidden Markov Model using JAX for numerical operations. This is a high-level API example.

```python
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
from dynamax.hidden_markov_model import GaussianHMM
```

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

### Initialize HMM Parameters and Sample Data

Source: https://github.com/probml/dynamax/blob/main/dynamax/hidden_markov_model/demos/low_rank_gaussian_hmm.ipynb

Sets up constants for the HMM, including the number of states, emission dimensionality, and covariance rank. It then initializes a LowRankGaussianHMM instance and samples true parameters, states, and emissions for a specified number of timesteps.

```python
# Global constants used to define the HMM
num_states = 3
emission_dim = 2
emission_cov_rank = 1

# Defined number of times to sample from the HMM
num_timesteps = 2000

# Set up an iterator to yield JAX PRNGKey's
keys = map(jr.PRNGKey, count())

lr_hmm = LowRankGaussianHMM(num_states, emission_dim, emission_cov_rank)
true_params, _ = lr_hmm.initialize(next(keys))
true_states, emissions = lr_hmm.sample(true_params, next(keys), num_timesteps)
```

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

### Initialize and Fit Gaussian HMM with EM

Source: https://github.com/probml/dynamax/blob/main/docs/index.rst

Initializes a Gaussian HMM using k-means and fits it using the Expectation-Maximization algorithm. Requires JAX keys and batch emissions data.

```python
params, props = hmm.initialize(key3, method="kmeans", emissions=batch_emissions)
params, lls = hmm.fit_em(params, props, batch_emissions, num_iters=20)
```

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

### Setup Plotting Style

Source: https://github.com/probml/dynamax/blob/main/docs/notebooks/hmm/autoregressive_hmm.ipynb

Configures the plotting style using Seaborn and defines a gradient colormap for visualizations. This helps in creating visually appealing plots.

```python
sns.set_style("white")

color_names = [
    "windows blue",
    "red",
    "amber",
    "faded green",
    "dusty purple",
    "orange",
    "brown",
    "pink"
]

colors = sns.xkcd_palette(color_names)
cmap = gradient_cmap(colors)
```

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

### Import necessary libraries for RBPF

Source: https://github.com/probml/dynamax/blob/main/docs/notebooks/slds/rbpf_maneuver.ipynb

Imports all required libraries and modules from dynamax, jax, tensorflow_probability, and other plotting utilities for the RBPF maneuver tracking example.

```python
import dynamax
import jax.numpy as jnp
import jax.random as jr
from functools import partial
import sys 
sys.path.append('/Users/kostastsampourakis/Desktop/code/Python/projects/dynamax')
from jax import vmap, jit
from jax.tree_util import tree_map
from dynamax.slds.inference import ParamsSLDS, LGParamsSLDS, DiscreteParamsSLDS, rbpfilter, rbpfilter_optimal
from dynamax.slds.models import SLDS
# import MVN from tfd
from tensorflow_probability.substrates.jax.distributions import MultivariateNormalFullCovariance as MVN

import seaborn as sns
import matplotlib.pyplot as plt
from functools import partial
from sklearn.preprocessing import OneHotEncoder
from jax.scipy.special import logit
from mpl_toolkits.mplot3d import Axes3D
import jax
```

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

### Initialize HMM Parameters

Source: https://github.com/probml/dynamax/blob/main/dynamax/hidden_markov_model/demos/bernoulli_hmm_example.ipynb

Sets the initial probabilities and transition matrix for the HMM, then initializes the model parameters. The true parameters are printed for inspection.

```python
initial_probs = jnp.ones((num_states,)) / (num_states * 1.0)
transition_matrix = 0.90 * jnp.eye(num_states) + 0.10 * jnp.ones((num_states, num_states)) / num_states
true_params, _ = true_hmm.initialize(key=jr.PRNGKey(0), initial_probs=initial_probs, transition_matrix=transition_matrix)
pprint(true_params)
```

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

### Color Initialization and Time Series Sampling

Source: https://github.com/probml/dynamax/blob/main/logo/make_logo.ipynb

Initializes colors for patches using a Gaussian process and samples subsequent timesteps as a random walk. This sets up the dynamic color changes for the logo animation.

```python
num_timesteps = 200

# Initialize the colors with a Gaussian process
centers = jnp.array([rect.xy for rect in patches])
length = 6 * RECT_PX
kernel = lambda xy1, xy2: jnp.exp(-0.5 * jnp.sum((xy1 - xy2)**2) / length**2)
cov = vmap(lambda xy2: vmap(lambda xy1: kernel(xy1, xy2))(centers))(centers)
cov += 1e-4 * jnp.eye(len(centers))

# Sample a time series of Gaussian processes, one for each primary color
fs = jnp.zeros((num_timesteps, num_patches, 3))

# Sample initial condition from mean zero GPs
seed = jr.PRNGKey(0)
key, seed = jr.split(seed)
fs = fs.at[0].set(tfd.MultivariateNormalFullCovariance(jnp.zeros(len(centers)), 3 * cov).sample(seed=key, sample_shape=(3,)).T)
fs = fs.at[0].add(-fs[0].mean(axis=0))

# Sample subsequent timesteps as random walk
for t in progress_bar(range(1, num_timesteps)):
    key, seed = jr.split(seed)
    fs = fs.at[t].set(tfd.MultivariateNormalFullCovariance(fs[t-1].T, 0.5 * cov).sample(seed=key).T)
    fs = fs.at[t].add(-fs[t].mean(axis=0))
    
# Sample primary and finecolors
colors = jnp.zeros((num_timesteps, num_patches), dtype=int)
key, seed = jr.split(seed)
colors = colors.at[0].set(tfd.Categorical(logits=jnp.repeat(fs[0], 3, axis=1)).sample(seed=key))
for t in progress_bar(range(1, num_timesteps)):
    key, seed = jr.split(seed)
    logits = jnp.repeat(fs[t], 3, axis=1) + 3 * one_hot(colors[t-1], 9)
    colors = colors.at[t].set(tfd.Categorical(logits=logits).sample(seed=key))
      
# Initialize the color of each rectangle
for rect, color in zip(patches, colors[0]):
    rect.set_facecolor(COLORS[color])
    rect.set_edgecolor('w')
    rect.set_linewidth(0.025)
```

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

### Initialize Linear Gaussian SSM parameters

Source: https://github.com/probml/dynamax/blob/main/docs/notebooks/linear_gaussian_ssm/kf_linreg.ipynb

Sets up the parameters for the Linear Gaussian State Space Model (LG-SSM) to represent the linear regression problem. This includes defining the state transition, emission matrices, initial state distribution, and observation noise covariance.

```python
F = jnp.eye(2)
Q = jnp.zeros((2, 2))  # No parameter drift.
obs_var = 1.0
R = jnp.ones((1, 1)) * obs_var
mu0 = jnp.zeros(2)
Sigma0 = jnp.eye(2) * 10.0


# the input_dim = 0 since we encode the covariates into the non-stationary emission matrix
lgssm = LinearGaussianSSM(state_dim = 2, emission_dim = 1, input_dim = 0)
params, _ = lgssm.initialize(
    initial_mean=mu0,
    initial_covariance=Sigma0,
    dynamics_weights=F,
    dynamics_covariance=Q,
    emission_weights=X[:, None, :], # (t, 1, D) where D = num input features
    emission_covariance=R,
    )
```

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

### Initialize Linear Gaussian State Space Model

Source: https://github.com/probml/dynamax/blob/main/docs/notebooks/linear_gaussian_ssm/kf_tracking.ipynb

Sets up a LinearGaussianSSM for tracking with specified state and emission dimensions. It defines and initializes the model parameters including initial mean and covariance, dynamics weights and covariance, and emission weights and covariance.

```python
state_dim = 4
emission_dim = 2
delta = 1.0

# Create object
lgssm = LinearGaussianSSM(state_dim, emission_dim)

# Manually chosen parameters
initial_mean = jnp.array([8.0, 10.0, 1.0, 0.0])
initial_covariance = jnp.eye(state_dim) * 0.1
dynamics_weights  = jnp.array([[1, 0, delta, 0],
                               [0, 1, 0, delta],
                               [0, 0, 1, 0],
                               [0, 0, 0, 1]])
dynamics_covariance = jnp.eye(state_dim) * 0.001
emission_weights = jnp.array([[1.0, 0, 0, 0],
                              [0, 1.0, 0, 0]])
emission_covariance = jnp.eye(emission_dim)

# Initialize model
params, _ = lgssm.initialize(jr.PRNGKey(0),
                             initial_mean=initial_mean,
                             initial_covariance=initial_covariance,
                             dynamics_weights=dynamics_weights,
                             dynamics_covariance=dynamics_covariance,
                             emission_weights=emission_weights,
                             emission_covariance=emission_covariance)



```

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

### Visualize Filtering Distribution with True States

Source: https://github.com/probml/dynamax/blob/main/docs/notebooks/hmm/casino_hmm_inference.ipynb

Calls the plotting function to display the filtering distribution. This example assumes 'posterior.filtered_probs' and 'true_states' are available from previous computations.

```python
plot_posterior_probs(posterior.filtered_probs, true_states,
                     title="Filtering Distribution")
```

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

### Import necessary libraries

Source: https://github.com/probml/dynamax/blob/main/docs/notebooks/linear_gaussian_ssm/kf_tracking.ipynb

Imports core JAX, Matplotlib, and dynamax modules required for setting up and visualizing the Linear Gaussian State Space Model.

```python
from jax import numpy as jnp
from jax import random as jr
from jax import vmap
from matplotlib import pyplot as plt

from dynamax.utils.plotting import plot_uncertainty_ellipses
from dynamax.linear_gaussian_ssm import LinearGaussianSSM
from dynamax.linear_gaussian_ssm import lgssm_smoother, lgssm_filter
```

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

### Defining Dataset Size Grid for 10 Classes

Source: https://github.com/probml/dynamax/blob/main/dynamax/generalized_gaussian_ssm/demos/cmgf_multiclass_logreg_demo.ipynb

Defines the grid of dataset sizes to be used for evaluation when working with 10 classes. This setup is typically used for scaling experiments.

```python
data_size_grid = list(range(150, 1000, 100)) + list(range(1000, 10000, 1000))
```

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

### Smooth and Forecast Emissions with Posterior Samples

Source: https://github.com/probml/dynamax/blob/main/docs/notebooks/linear_gaussian_ssm/lgssm_hmc.ipynb

Approximates the variance of smoothed and forecast emissions using the law of total variance from posterior samples. Requires prior setup of model parameters and emissions.

```python
(smooth_emissions_samples, smooth_emissions_std_samples, 
forecast_emissions_samples, forecast_emissions_std_samples) =
    vmap(lambda params: smooth_and_forecast(model, params, emissions))(param_samples)

# Compute the mean and standard deviation of the samples
burnin = 200
est_smooth_emissions = jnp.mean(smooth_emissions_samples[burnin:], axis=0)
est_forecast_emissions = jnp.mean(forecast_emissions_samples[burnin:], axis=0)
est_smooth_emissions_std = jnp.sqrt(
    jnp.mean(smooth_emissions_std_samples[burnin:] ** 2, axis=0) + 
    jnp.var(smooth_emissions_samples[burnin:], axis=0))
est_forecast_emissions_std = jnp.sqrt(
    jnp.mean(forecast_emissions_std_samples[burnin:] ** 2, axis=0) + 
    jnp.var(forecast_emissions_samples[burnin:], axis=0))

```

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

### Get MLP Flattened Parameters and Apply Function

Source: https://github.com/probml/dynamax/blob/main/docs/notebooks/generalized_gaussian_ssm/cmgf_mlp_classification_demo.ipynb

Initializes an MLP model, flattens its parameters, and creates an apply function. This is useful for integrating with optimization or inference frameworks that require flattened parameters.

```python
def get_mlp_flattened_params(model_dims, key=0):
    if isinstance(key, int):
        key = jr.PRNGKey(key)

    # Define MLP model
    input_dim, features = model_dims[0], model_dims[1:]
    model = MLP(features)
    dummy_input = jnp.ones((input_dim,))

    # Initialize parameters using dummy input
    params = model.init(key, dummy_input)
    flat_params, unflatten_fn = ravel_pytree(params)

    # Define apply function
    def apply(flat_params, x, model, unflatten_fn):
        return model.apply(unflatten_fn(flat_params), jnp.atleast_1d(x))

    apply_fn = partial(apply, model=model, unflatten_fn=unflatten_fn)

    return model, flat_params, unflatten_fn, apply_fn
```

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

### Import inference and parameter utilities

Source: https://github.com/probml/dynamax/blob/main/dynamax/generalized_gaussian_ssm/demos/dirichlet_kalman-filter_demo.ipynb

Imports necessary components for Kalman filtering inference and parameter creation from the dynamax library. These are essential for setting up and running the SSM.

```python
from sklearn.base import BaseEstimator, ClassifierMixin
from sklearn.metrics import log_loss

from dynamax.linear_gaussian_ssm.inference import ParamsLGSSM, lgssm_filter, make_lgssm_params
```

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

### Import necessary libraries for LG-SSM inference

Source: https://github.com/probml/dynamax/blob/main/docs/notebooks/linear_gaussian_ssm/lgssm_parallel_inference.ipynb

Imports JAX, NumPy, random number generation, plotting tools, and specific functions for LG-SSM filtering and smoothing from the dynamax library. This setup is required for both sequential and parallel inference.

```python
import jax
from jax import numpy as jnp
from jax import random as jr
from jax import block_until_ready
from matplotlib import pyplot as plt

from dynamax.linear_gaussian_ssm import lgssm_smoother, parallel_lgssm_smoother
from dynamax.linear_gaussian_ssm import LinearGaussianSSM
```

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

### Import JAX and Utilities

Source: https://github.com/probml/dynamax/blob/main/docs/notebooks/generalized_gaussian_ssm/cmgf_mlp_classification_demo.ipynb

Imports core JAX libraries and utilities for numerical computation, random number generation, and plotting.

```python
from typing import Sequence
from functools import partial

import matplotlib.pyplot as plt
import matplotlib.cm as cm
import jax
import jax.numpy as jnp
import jax.random as jr
from jax.flatten_util import ravel_pytree
```

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

### Sample and Fit Gaussian HMM

Source: https://github.com/probml/dynamax/blob/main/docs/index.rst

Demonstrates sampling data from a Gaussian HMM and then fitting the model parameters using EM. Requires JAX and matplotlib.

```python
import dynamax.hidden_markov_model as hmm
import jax.random as jr
import matplotlib.pyplot as plt

key1, key2, key3 = jr.split(jr.PRNGKey(0), 3)
num_states = 3
emission_dim = 2
num_timesteps = 1000

# Make a Gaussian HMM and sample data from it
hmm = hmm.GaussianHMM(num_states, emission_dim)
true_params, _ = hmm.initialize(key1)
true_states, emissions = hmm.sample(true_params, key2, num_timesteps)

# Make a new Gaussian HMM and fit it with EM
params, props = hmm.initialize(key3, method="kmeans", emissions=emissions)
params, lls = hmm.fit_em(params, props, emissions, num_iters=20)

# Plot the marginal log probs across EM iterations
plt.plot(lls)
plt.xlabel("EM iterations")
plt.ylabel("marginal log prob.")

# Use fitted model for posterior inference
post = hmm.smoother(params, emissions)
print(post.smoothed_probs.shape) # (1000, 3)
```

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

### MLP Initialization and Parameter Handling

Source: https://github.com/probml/dynamax/blob/main/docs/notebooks/nonlinear_gaussian_ssm/ekf_mlp.ipynb

Generates an MLP model, initializes its parameters using dummy input, and provides functions for parameter flattening, unflattening, and applying the model. This is useful for setting up an MLP for training or inference.

```python
def get_mlp_flattened_params(model_dims, key=0):
    """Generate MLP model, initialize it using dummy input, and
    return the model, its flattened initial parameters, function
    to unflatten parameters, and apply function for the model.
    Args:
        model_dims (List): List of [input_dim, hidden_dim, ..., output_dim]
        key (PRNGKey): Random key. Defaults to 0.
    Returns:
        model: MLP model with given feature dimensions.
        flat_params: Flattened parameters initialized using dummy input.
        unflatten_fn: Function to unflatten parameters.
        apply_fn: fn(flat_params, x) that returns the result of applying the model.
    """
    if isinstance(key, int):
        key = jr.PRNGKey(key)

    # Define MLP model
    input_dim, features = model_dims[0], model_dims[1:]
    model = MLP(features)
    dummy_input = jnp.ones((input_dim,))

    # Initialize parameters using dummy input
    params = model.init(key, dummy_input)
    flat_params, unflatten_fn = ravel_pytree(params)

    # Define apply function
    def apply(flat_params, x, model, unflatten_fn):
        return model.apply(unflatten_fn(flat_params), jnp.atleast_1d(x))

    apply_fn = partial(apply, model=model, unflatten_fn=unflatten_fn)

    return model, flat_params, unflatten_fn, apply_fn

```

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

### Prepare Input for Emission Matrix

Source: https://github.com/probml/dynamax/blob/main/dynamax/generalized_gaussian_ssm/demos/dirichlet_kalman-filter_demo.ipynb

Prepares the input data by adding a bias term and constructing the emission matrix using Kronecker products. This is a preparatory step for the Kalman filter.

```python
X_plus_bias = jnp.hstack((jnp.ones((len(X),1)),X))
input_dim = X_plus_bias.shape[-1]

emission_matrix = jnp.stack([jnp.kron(jnp.eye(input_dim), row) for row in X_plus_bias])
emission_matrix[0]
```

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

### Initialize and Run Unscented Kalman Filter

Source: https://github.com/probml/dynamax/blob/main/docs/notebooks/nonlinear_gaussian_ssm/ekf_ukf_spiral.ipynb

Initializes UKF hyperparameters and applies the UKF to the given parameters and emissions to obtain filtered means and covariances. Requires prior parameters, emissions data, and UKF hyperparameters.

```python
hyperparams = UKFHyperParams()
ukf_post = ukf(params, emissions, hyperparams)
ukf_means, ukf_covs = ukf_post.filtered_means, ukf_post.filtered_covariances
```

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

### Import necessary libraries

Source: https://github.com/probml/dynamax/blob/main/docs/notebooks/nonlinear_gaussian_ssm/ekf_ukf_spiral.ipynb

Imports JAX, Matplotlib, and specific modules from dynamax for nonlinear Gaussian SSMs, EKF, and UKF.

```python
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt

from dynamax.utils.plotting import  plot_uncertainty_ellipses
from dynamax.nonlinear_gaussian_ssm import ParamsNLGSSM, NonlinearGaussianSSM, UKFHyperParams
from dynamax.nonlinear_gaussian_ssm import extended_kalman_filter as ekf
from dynamax.nonlinear_gaussian_ssm import unscented_kalman_filter as ukf
```

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

### Initialize EKF Parameters for Pendulum Model

Source: https://github.com/probml/dynamax/blob/main/docs/notebooks/nonlinear_gaussian_ssm/ekf_ukf_pendulum.ipynb

Sets up the parameters for the Extended Kalman Filter using the Pendulum model. This includes initial state mean and covariance, dynamics and emission functions, and their respective covariances. Requires `PendulumParams` and `ParamsNLGSSM` from Dynamax.

```python
pendulum_params = PendulumParams()

# Define parameters for EKF
ekf_params = ParamsNLGSSM(
    initial_mean=pendulum_params.initial_state,
    initial_covariance=jnp.eye(states.shape[-1]) * 0.1,
    dynamics_function=pendulum_params.dynamics_function,
    dynamics_covariance=pendulum_params.dynamics_covariance,
    emission_function=pendulum_params.emission_function,
    emission_covariance=pendulum_params.emission_covariance,
)

ekf_posterior = ekf(ekf_params, obs)
```

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

### Initialize and Fit HMM with Exact EM

Source: https://github.com/probml/dynamax/blob/main/dynamax/hidden_markov_model/demos/poisson_hmm_neurons.ipynb

Initializes HMM parameters and fits the model to emissions using exact EM updates. This is useful for training the HMM when the emission probabilities can be computed exactly.

```python
params, param_props = hmm.initialize(jr.PRNGKey(1234))
test_hmm, lps = hmm.fit_em(params, param_props, emissions)
```

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

### Generate Synthetic Data and Initialize Parameters

Source: https://github.com/probml/dynamax/blob/main/dynamax/linear_gaussian_ssm/demos/lgssm_blocked_gibbs.ipynb

Generates synthetic data from a true LGSSM and initializes parameters for a test model using Maximum A Posteriori (MAP) estimation via EM. This snippet sets up the data and initial parameters for subsequent inference.

```python
from itertools import count
state_dim = 2
emission_dim = 10
num_timesteps = 100

keys = map(jr.PRNGKey, count())

true_model = LinearGaussianConjugateSSM(state_dim, emission_dim)
true_params, _ = true_model.initialize(next(keys))

true_states, emissions = true_model.sample(true_params, next(keys), num_timesteps)

num_iters = 100
test_model = LinearGaussianConjugateSSM(state_dim, emission_dim)
test_params, test_param_props = test_model.initialize(next(keys))
test_params, marginal_lls = test_model.fit_em(test_params, test_param_props, emissions, num_iters=num_iters)

# Compute predicted emissions
emission_weights = test_params.emissions.weights
emission_bias = test_params.emissions.bias
emission_cov = test_params.emissions.cov
posterior = test_model.smoother(test_params, emissions)
smoothed_emissions = posterior.smoothed_means @ emission_weights.T + emission_bias
smoothed_emissions_cov = emission_weights @ posterior.smoothed_covariances @ emission_weights.T + emission_cov
smoothed_emissions_std = jnp.sqrt(jnp.array([smoothed_emissions_cov[:, i, i] for i in range(emission_dim)]))

spc = 3
plt.figure(figsize=(10, 4))
for i in range(emission_dim):
    plt.plot(emissions[:, i] + spc * i, "--k", label="observed" if i == 0 else None)
    ln = plt.plot(smoothed_emissions[:, i] + spc * i,
                  label="smoothed" if i == 0 else None)[0]
    plt.fill_between(
        jnp.arange(num_timesteps),
        spc * i + smoothed_emissions[:, i] - 2 * jnp.sqrt(smoothed_emissions_std[i]),
        spc * i + smoothed_emissions[:, i] + 2 * jnp.sqrt(smoothed_emissions_std[i]),
        color=ln.get_color(),
        alpha=0.25,
    )
plt.xlabel("time")
plt.xlim(0, num_timesteps - 1)
plt.ylabel("true and predicted emissions")
plt.legend()
```

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

### Import necessary libraries for Poisson LDS demo

Source: https://github.com/probml/dynamax/blob/main/docs/notebooks/generalized_gaussian_ssm/cmgf_poisson_demo.ipynb

Imports core JAX libraries, plotting tools, Poisson distribution, and specific modules from dynamax for Generalized Gaussian State Space Models (GGSSM) and the Conditional Moments Gaussian Smoother (CMGS).

```python
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt

from jax import vmap
from jax.nn import softplus
from tensorflow_probability.substrates.jax.distributions import Poisson

from dynamax.generalized_gaussian_ssm import ParamsGGSSM, GeneralizedGaussianSSM, EKFIntegrals
from dynamax.generalized_gaussian_ssm import conditional_moments_gaussian_smoother as cmgs
from dynamax.utils.utils import random_rotation
```

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

### Import necessary libraries

Source: https://github.com/probml/dynamax/blob/main/docs/notebooks/linear_gaussian_ssm/lgssm_learning.ipynb

Imports core JAX, Matplotlib, and dynamax modules required for LG-SSM parameter estimation and state inference.

```python
import jax.numpy as jnp
import jax.random as jr

from jax import vmap
from matplotlib import pyplot as plt
from optax import adam

from dynamax.linear_gaussian_ssm import LinearGaussianSSM
from dynamax.utils.utils import monotonically_increasing, random_rotation
```

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

### Import necessary libraries

Source: https://github.com/probml/dynamax/blob/main/dynamax/hidden_markov_model/demos/bernoulli_hmm_example.ipynb

Imports core libraries including JAX for numerical computation, matplotlib for plotting, and the BernoulliHMM model from dynamax.

```python
from functools import partial
import jax.random as jr
import jax.numpy as jnp
from jax import vmap
import matplotlib.pyplot as plt
from pprint import pprint

from dynamax.hidden_markov_model import BernoulliHMM
```

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

### Initialize and Fit Poisson GLM-HMM

Source: https://github.com/probml/dynamax/blob/main/docs/notebooks/hmm/custom_hmm.ipynb

Initializes a Poisson GLM-HMM model and fits it to emission and input data using the EM algorithm. Requires JAX and dynamax libraries.

```python
init_key = jr.PRNGKey(12345)
test_num_states = num_states

# Initialize with new random seed
model = PoissonGLMHMM(num_states=test_num_states, emission_dim=emission_dim, input_dim=input_dim)
params, props = model.initialize(key=init_key)

# Fit with EM
fitted_params, lps = model.fit_em(params=params, props=props, emissions=emissions, inputs=inputs)
```

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

### Check available JAX devices and set CPU mode

Source: https://github.com/probml/dynamax/blob/main/docs/notebooks/linear_gaussian_ssm/lgssm_parallel_inference.ipynb

Prints the available JAX devices and their platforms. It sets a `cpu_mode` flag based on whether the primary device is a CPU, which might influence subsequent computations or configurations.

```python
print(jax.devices())
print(jax.devices()[0].platform)
if jax.devices()[0].platform == 'cpu':
    cpu_mode = True
else:
    cpu_mode = False
```

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

### Create a Bernoulli HMM

Source: https://github.com/probml/dynamax/blob/main/dynamax/hidden_markov_model/demos/bernoulli_hmm_example.ipynb

Initializes a Bernoulli HMM with specified number of states and emission dimensions. Prior concentrations for emissions are set to 1.0.

```python
num_states = 5
emission_dim = 10
true_hmm = BernoulliHMM(num_states, emission_dim, 
                        emission_prior_concentration0=1.0, 
                        emission_prior_concentration1=1.0)
```

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

### Import necessary libraries

Source: https://github.com/probml/dynamax/blob/main/docs/notebooks/linear_gaussian_ssm/kf_linreg.ipynb

Imports the required components from JAX, Matplotlib, and dynamax for implementing the Kalman filter for linear regression.

```python
from jax import numpy as jnp
from jax import vmap
from matplotlib import pyplot as plt
from dynamax.linear_gaussian_ssm import LinearGaussianSSM
```

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

### Import CMGF Components

Source: https://github.com/probml/dynamax/blob/main/docs/notebooks/generalized_gaussian_ssm/cmgf_mlp_classification_demo.ipynb

Imports necessary components for Generalized Gaussian SSM and CMGF from the dynamax library.

```python
from dynamax.generalized_gaussian_ssm import ParamsGGSSM, EKFIntegrals
from dynamax.generalized_gaussian_ssm import conditional_moments_gaussian_filter
```

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

### Import Libraries

Source: https://github.com/probml/dynamax/blob/main/dynamax/hidden_markov_model/demos/poisson_hmm_changepoint.ipynb

Imports essential libraries for JAX, TensorFlow Probability, Optax, and Dynamax, including specific modules for HMMs and Poisson distributions.

```python
import jax.numpy as jnp
import jax.random as jr
from jax import jit, lax, vmap, value_and_grad

import tensorflow_probability.substrates.jax as tfp
tfd = tfp.distributions
tfb = tfp.bijectors

import optax
from itertools import count
from functools import partial
from matplotlib import pylab as plt

from dynamax.hidden_markov_model import PoissonHMM, HMM
from dynamax.parameters import ParameterProperties
```

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

### Initialize LGSSM and Freeze Parameters for HMC

Source: https://github.com/probml/dynamax/blob/main/docs/notebooks/linear_gaussian_ssm/lgssm_hmc.ipynb

Initializes a LinearGaussianConjugateSSM model and freezes specific parameters (transition and initial parameters) to learn only covariance matrices. Requires PRNG keys and true parameter values.

```python
# Freeze transition parameters and initial parameters, so that only covariance matrices are learned
init_key, hmc_key = jr.split(jr.PRNGKey(42))
model = LinearGaussianConjugateSSM(state_dim, emission_dim)
params, param_props = model.initialize(
    init_key,
    emission_weights=true_params.emissions.weights,
    emission_bias=true_params.emissions.bias
)

# Set transition parameters and initial parameters to true values and mark as frozen
param_props.emissions.weights.trainable = False
param_props.emissions.bias.trainable = False

```

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

### Initialize and Sample LGSSM Data

Source: https://github.com/probml/dynamax/blob/main/docs/notebooks/linear_gaussian_ssm/lgssm_learning.ipynb

Constructs a LinearGaussianSSM with randomly initialized parameters, including dynamics matrix and covariance. It then samples latent states and emissions from this model. Ensure JAX and Dynamax are imported as 'jnp' and 'jr' respectively.

```python
state_dim = 2
emission_dim = 10
num_timesteps = 100
k1, k2, k3 = jr.split(jr.PRNGKey(0), 3)

# Construct the true model with randomly initialized parameters
true_A = 0.99 * random_rotation(seed=k1, n=state_dim, theta=jnp.pi / 10)
true_Sigma = 0.01 * jnp.eye(state_dim)
true_model = LinearGaussianSSM(state_dim, emission_dim)
true_params, param_props = true_model.initialize(
    key=k1, dynamics_weights=true_A, dynamics_covariance=true_Sigma)
            
# Sample states and emissions from the true model
true_states, emissions = true_model.sample(true_params, k3, num_timesteps)
```

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

### Import necessary libraries

Source: https://github.com/probml/dynamax/blob/main/docs/notebooks/generalized_gaussian_ssm/cmgf_logistic_regression_demo.ipynb

Imports core libraries for plotting, JAX operations, and dynamax components for generalized Gaussian SSMs.

```python
import matplotlib.pyplot as plt
import jax
import jax.numpy as jnp
import jax.random as jr

from jax.scipy.optimize import minimize

from dynamax.generalized_gaussian_ssm import ParamsGGSSM, EKFIntegrals, UKFIntegrals, GHKFIntegrals
from dynamax.generalized_gaussian_ssm import conditional_moments_gaussian_filter
```

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

### Initialize Categorical HMM Parameters

Source: https://github.com/probml/dynamax/blob/main/docs/notebooks/hmm/casino_hmm_learning.ipynb

Defines the number of states, emissions, and classes, then sets the initial probabilities, transition matrix, and emission probabilities for a CategoricalHMM. This snippet is used to set up the ground truth parameters for data generation.

```python
num_states = 2      # two types of dice (fair and loaded)
num_emissions = 1   # only one die is rolled at a time
num_classes = 6     # each die has six faces

initial_probs = jnp.array([0.5, 0.5])
transition_matrix = jnp.array([[0.95, 0.05], 
                               [0.10, 0.90]])
emission_probs = jnp.array([[1/6,  1/6,  1/6,  1/6,  1/6,  1/6],    # fair die
                            [1/10, 1/10, 1/10, 1/10, 1/10, 5/10]])  # loaded die


# Construct the HMM
hmm = CategoricalHMM(num_states, num_emissions, num_classes)

# Initialize the parameters struct with known values
params, _ = hmm.initialize(initial_probs=initial_probs,
                           transition_matrix=transition_matrix,
                           emission_probs=emission_probs.reshape(num_states, num_emissions, num_classes))
```

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

### Initialize Poisson HMM

Source: https://github.com/probml/dynamax/blob/main/dynamax/hidden_markov_model/demos/poisson_hmm_neurons.ipynb

Initializes a Poisson Hidden Markov Model with specified number of states and neurons, and sets up initial probabilities and transition matrix for the true model.

```python
num_states = 5
num_neurons = 10
hmm = PoissonHMM(num_states, num_neurons, emission_prior_concentration=3.0, emission_prior_rate=1.0)

initial_probs = jnp.ones((num_states,)) / num_states
transition_matrix = 0.90 * jnp.eye(num_states) + 0.10 * jnp.ones((num_states, num_states)) / num_states
true_params, param_props = hmm.initialize(jr.PRNGKey(0), initial_probs=initial_probs, transition_matrix=transition_matrix)
```

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

### Sample Data from Gaussian HMM

Source: https://github.com/probml/dynamax/blob/main/docs/notebooks/hmm/gaussian_hmm.ipynb

Generates training and test data from a Gaussian HMM. Requires defining HMM parameters and using JAX for sampling.

```python
num_train_batches = 3
num_test_batches = 1
num_timesteps = 100

# Make an HMM and sample data and true underlying states
true_num_states = 5
emission_dim = 2
hmm = GaussianHMM(true_num_states, emission_dim)

# Specify parameters of the HMM
initial_probs = jnp.ones(true_num_states) / true_num_states
transition_matrix = 0.80 * jnp.eye(true_num_states) \
    + 0.15 * jnp.roll(jnp.eye(true_num_states), 1, axis=1) \
    + 0.05 / true_num_states
emission_means = jnp.column_stack([
    jnp.cos(jnp.linspace(0, 2 * jnp.pi, true_num_states + 1))[:-1],
    jnp.sin(jnp.linspace(0, 2 * jnp.pi, true_num_states + 1))[:-1],
    jnp.zeros((true_num_states, emission_dim - 2)),
    ])
emission_covs = jnp.tile(0.1**2 * jnp.eye(emission_dim), (true_num_states, 1, 1))
        
true_params, _ = hmm.initialize(initial_probs=initial_probs,
                                transition_matrix=transition_matrix,
                                emission_means=emission_means,
                                emission_covariances=emission_covs)

# Sample train, validation, and test data
train_key, val_key, test_key = jr.split(jr.PRNGKey(0), 3)
f = vmap(partial(hmm.sample, true_params, num_timesteps=num_timesteps))
train_true_states, train_emissions = f(jr.split(train_key, num_train_batches))
test_true_states,  test_emissions  = f(jr.split(test_key, num_test_batches))
```

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

### Prepare Test Sequences

Source: https://github.com/probml/dynamax/blob/main/dynamax/hidden_markov_model/demos/bach_chorales_hmm.ipynb

Fetches test data and filters sequences to include only present notes. Assumes 'present_notes' is defined.

```python
test_lengths, test_sequences = fetch()

# find all the notes that are present at least once in the training set
test_sequences = test_sequences[..., present_notes]
```

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

### Import Necessary Libraries for HMM Learning

Source: https://github.com/probml/dynamax/blob/main/docs/notebooks/hmm/casino_hmm_learning.ipynb

Imports core libraries including JAX for numerical computation, Optax for optimization, Matplotlib for plotting, and specific components from Dynamax for HMM modeling.

```python
from functools import partial

import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt
import optax
from jax import vmap

from dynamax.hidden_markov_model import CategoricalHMM
```