From f85ff15a4101293c363e133ce78148985c09aca1 Mon Sep 17 00:00:00 2001
From: Andy Maloney <amaloney@quansight.com>
Date: Wed, 21 Sep 2022 17:40:37 -0500
Subject: [PATCH 1/7] Marginal diagnostic tool

This commit includes the marginal 1D diagnostic tool with JavaScript
callbacks.
---
 .../ppl/diagnostics/tools/js/package.json     | 25 ----------
 .../ppl/diagnostics/tools/utils/base.py       | 46 +++++++++++++++++++
 .../tools/utils/model_serializers.py          | 31 ++++++++++---
 3 files changed, 71 insertions(+), 31 deletions(-)
 create mode 100644 src/beanmachine/ppl/diagnostics/tools/utils/base.py

diff --git a/src/beanmachine/ppl/diagnostics/tools/js/package.json b/src/beanmachine/ppl/diagnostics/tools/js/package.json
index 38e7770cb5..7518ebe720 100644
--- a/src/beanmachine/ppl/diagnostics/tools/js/package.json
+++ b/src/beanmachine/ppl/diagnostics/tools/js/package.json
@@ -13,8 +13,6 @@
     "fast-kde": "^0.2.1"
   },
   "devDependencies": {
-    "@types/node": "^18.0.4",
-    "@typescript-eslint/eslint-plugin": "^5.30.5",
     "@typescript-eslint/parser": "^5.30.5",
     "eslint": "^8.19.0",
     "eslint-config-airbnb": "^19.0.4",
@@ -24,32 +22,9 @@
     "eslint-plugin-prefer-arrow": "^1.2.3",
     "eslint-plugin-react": "^7.28.0",
     "eslint-plugin-react-hooks": "^4.3.0",
-    "prettier": "^2.7.1",
     "ts-loader": "^9.3.1",
-    "ts-node": "^10.9.1",
     "typescript": "^4.7.4",
     "webpack": "^5.74.0",
     "webpack-cli": "^4.10.0"
-  },
-  "overrides": {
-    "cwise": "$cwise",
-    "minimist": "$minimist",
-    "quote-stream": "$quote-stream",
-    "static-eval": "$static-eval",
-    "static-module": "$static-module",
-    "typedarray-pool": "$typedarray-pool"
-  },
-  "peerDependencies": {
-    "@types/cwise": "^1.0.4",
-    "@types/minimist": "^1.2.2",
-    "@types/static-eval": "^0.2.31",
-    "@types/typedarray-pool": "^1.1.2",
-    "buffer": "^6.0.3",
-    "cwise": "^1.0.10",
-    "minimist": "^1.2.6",
-    "quote-stream": "^1.0.2",
-    "static-eval": "2.1.0",
-    "static-module": "^3.0.4",
-    "typedarray-pool": "^1.2.0"
   }
 }
diff --git a/src/beanmachine/ppl/diagnostics/tools/utils/base.py b/src/beanmachine/ppl/diagnostics/tools/utils/base.py
new file mode 100644
index 0000000000..8f290ec068
--- /dev/null
+++ b/src/beanmachine/ppl/diagnostics/tools/utils/base.py
@@ -0,0 +1,46 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Base class for diagnostic tools of a Bean Machine model."""
+
+import re
+from abc import ABC, abstractmethod
+from typing import TypeVar
+
+from beanmachine.ppl.diagnostics.tools import JS_DIST_DIR
+from beanmachine.ppl.diagnostics.tools.utils import plotting_utils
+from beanmachine.ppl.diagnostics.tools.utils.model_serializers import serialize_bm
+from beanmachine.ppl.inference.monte_carlo_samples import MonteCarloSamples
+from IPython.display import display, HTML
+
+
+T = TypeVar("T", bound="Base")
+
+
+class Base(ABC):
+    @abstractmethod
+    def __init__(self: T, mcs: MonteCarloSamples) -> None:
+        self.data = serialize_bm(mcs)
+        self.rv_names = list(self.data.keys())
+        self.num_chains = mcs.num_chains
+        self.num_draws = mcs.get_num_samples()
+        self.palette = plotting_utils.choose_palette(self.num_chains)
+        self.js = self.load_js()
+
+    def load_js(self: T) -> str:
+        name = self.__class__.__name__
+        name_tokens = re.findall(r"[A-Z][^A-Z]*", name)
+        name = "_".join(name_tokens)
+        path = JS_DIST_DIR.joinpath(f"{name.lower()}.js")
+        with path.open() as f:
+            js = f.read()
+        return js
+
+    def show(self: T):
+        display(HTML(self.create_document()))
+
+    @abstractmethod
+    def create_document(self: T) -> str:
+        ...
diff --git a/src/beanmachine/ppl/diagnostics/tools/utils/model_serializers.py b/src/beanmachine/ppl/diagnostics/tools/utils/model_serializers.py
index 109b588ea0..67c2b418e9 100644
--- a/src/beanmachine/ppl/diagnostics/tools/utils/model_serializers.py
+++ b/src/beanmachine/ppl/diagnostics/tools/utils/model_serializers.py
@@ -21,10 +21,29 @@ def serialize_bm(samples: MonteCarloSamples) -> Dict[str, List[List[float]]]:
         Dict[str, List[List[float]]]: The JSON serializable object for use in the
             diagnostics tools.
     """
-    model = dict(
-        sorted(
-            {str(key): value.tolist() for key, value in samples.items()}.items(),
-            key=lambda item: item[0],
-        ),
-    )
+    rv_identifiers = list(samples.keys())
+    reshaped_data = {}
+    for rv_identifier in rv_identifiers:
+        rv_data = samples[rv_identifier]
+        rv_shape = rv_data.shape
+        num_rv_chains = rv_shape[0]
+        for rv_chain in range(num_rv_chains):
+            chain_data = rv_data[rv_chain, :]
+            chain_shape = chain_data.shape
+            if len(chain_shape) > 3 and 1 not in list(chain_shape):
+                msg = (
+                    "Unable to handle data with dimensionality larger than " "mxnxkx1."
+                )
+                raise ValueError(msg)
+            elif len(chain_shape) == 3 and 1 in list(chain_shape):
+                if chain_shape[1] == 1 in list(chain_shape):
+                    reshape_dimensions = chain_shape[2]
+                else:
+                    reshape_dimensions = chain_shape[1]
+                for i, reshape_dimension in enumerate(range(reshape_dimensions)):
+                    data = rv_data[rv_chain, :, reshape_dimension].reshape(-1)
+                    reshaped_data[f"{str(rv_identifier)}[{i}]"] = data.tolist()
+            elif len(chain_shape) == 1:
+                reshaped_data[f"{str(rv_identifier)}"] = rv_data[rv_chain, :].tolist()
+    model = dict(sorted(reshaped_data.items(), key=lambda item: item[0]))
     return model

From 37c00676cf63622c56544cf62d45d805cdcead15 Mon Sep 17 00:00:00 2001
From: Andy Maloney <amaloney@quansight.com>
Date: Thu, 22 Sep 2022 16:35:11 -0500
Subject: [PATCH 2/7] Documentation & remove arviz

- Removes ArviZ as a dependency from the diagnostic tool.
- Adds jsdoc Documentation for the TypeScript methods.
- Fixes flake8 linting error.
- Adds clarifying comments to the tool.
- Adds missing copyright notice.
---
 .../ppl/diagnostics/tools/marginal1d/utils.py |  3 +-
 .../ppl/diagnostics/tools/utils/hdi.py        | 90 +++++++++++++++++++
 2 files changed, 92 insertions(+), 1 deletion(-)
 create mode 100644 src/beanmachine/ppl/diagnostics/tools/utils/hdi.py

diff --git a/src/beanmachine/ppl/diagnostics/tools/marginal1d/utils.py b/src/beanmachine/ppl/diagnostics/tools/marginal1d/utils.py
index 6ef4dfebf7..2061c890ea 100644
--- a/src/beanmachine/ppl/diagnostics/tools/marginal1d/utils.py
+++ b/src/beanmachine/ppl/diagnostics/tools/marginal1d/utils.py
@@ -8,8 +8,9 @@
 from typing import List
 
 import numpy as np
+import statsmodels.api as sm
 from beanmachine.ppl.diagnostics.tools.marginal1d import typing
-from beanmachine.ppl.diagnostics.tools.utils import plotting_utils
+from beanmachine.ppl.diagnostics.tools.utils import hdi, plotting_utils
 from bokeh.models.annotations import Band, LabelSet
 from bokeh.models.glyphs import Circle, Line
 from bokeh.models.layouts import Column, Row
diff --git a/src/beanmachine/ppl/diagnostics/tools/utils/hdi.py b/src/beanmachine/ppl/diagnostics/tools/utils/hdi.py
new file mode 100644
index 0000000000..d14594beff
--- /dev/null
+++ b/src/beanmachine/ppl/diagnostics/tools/utils/hdi.py
@@ -0,0 +1,90 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+from typing import Dict, List, Union
+
+
+def interval(
+    data: List[float],
+    hdi_probability: float,
+) -> Dict[str, Union[List[float], float]]:
+    """Find the lower and upper bound of the Highest Density Interval (HDI).
+
+    Parameters
+    ----------
+    data : List[float]
+        Raw random variable data from the model.
+    hdi_probability : float
+        The highest density interval probability to use when calculating the HDI.
+
+    Returns
+    -------
+    Dict[str, Union[List[float], float]]
+        A dictionary defining the lower and upper bound for the HDI of the given data,
+        as well as the indices of the lower and upper bound for the sorted raw data.
+    """
+    N = len(data)
+    sorted_data = sorted(data)
+    stop_index = math.floor(hdi_probability * N)
+    start_index = N - stop_index
+    left_data = sorted_data[stop_index:]
+    right_data = sorted_data[:start_index]
+    hdi = []
+    for i in range(len(left_data)):
+        hdi.append(left_data[i] - right_data[i])
+    lower_index = hdi.index(min(hdi))
+    upper_index = lower_index + stop_index
+    return {
+        "lower_bound": sorted_data[lower_index],
+        "upper_bound": sorted_data[upper_index],
+        "lower_bound_index": lower_index,
+        "upper_bound_index": upper_index,
+    }
+
+
+def data(
+    rv_data: List[float],
+    marginal_x: List[float],
+    marginal_y: List[float],
+    hdi_probability: float,
+) -> Dict[str, Union[List[float], float]]:
+    """Construct x and y arrays from the HDI bounds using the raw random variable dataself.
+
+    It also returns the actual HDI bounds of the random variable data.
+
+    Parameters
+    ----------
+    rv_data : List[float]
+        Raw random variable data from the model.
+    marginalX : List[float]
+        The support of the Kernel Density Estimate of the random variable.
+    marginalY : List[float]
+        The Kernel Density Estimate of the random variable.
+    hdiProbability : float
+    The highest density interval probability to use when alculating the HDI.
+
+    Returns
+    -------
+    Dict[str, Union[List[float], float]]
+        The x array for the HDI data is the "lower" key, while the y array for the HDI
+        data is the "upper" key and the "base" key is the support along the axis the HDI
+        is calculated. This nomenclature mirrors the Bokeh Band object, which is the
+        annotation type we will use when rendering the HDI intervals.
+    """
+    hdi = interval(rv_data, hdi_probability)
+    base = []
+    upper = []
+    for i in range(len(marginal_x)):
+        if marginal_x[i] <= hdi["upper_bound"] and marginal_x[i] >= hdi["lower_bound"]:
+            base.append(marginal_x[i])
+            upper.append(marginal_y[i])
+    return {
+        "base": base,
+        "lower": [0] * len(base),
+        "upper": upper,
+        "lower_bound": hdi["lower_bound"],
+        "upper_bound": hdi["upper_bound"],
+    }

From 27b8c696d94e061bcfe28ef8337c4f4bc5332f26 Mon Sep 17 00:00:00 2001
From: Andy Maloney <amaloney@quansight.com>
Date: Fri, 23 Sep 2022 16:28:26 -0500
Subject: [PATCH 3/7] Refactor

- Aligns type names between Python and TypeScript. There is now a
  one-to-one correspondence between Python TypedDict objects and the
  TypeScript interface objects.
- Redundancy in the Python TypedDict objects has been removed.
- Fixes tooltips not rendering correctly.
- Updates the link in the Help Tab of the tool such that when clicked it
  opens a new tab in the browser.
- Moves a few peerDependencies to the devDependencies section in the
  package.json file. This was done because yarn issued a warning
  indicating there were missing peer dependencies.
- Refactors the tool such that no calculations for the data displayed is
  done in Python. All updates and all calculations for the tool data is
  done exclusively in the browser. The rational for this change was to
  prevent users from having initial data calculated by Python, and
  subsequent updates calculated by JavaScript in the browser when
  interacting with the tool. There are notes in the code indicating why
  this change was made.
---
 .../ppl/diagnostics/tools/js/yarn.lock        | 10 +++
 .../ppl/diagnostics/tools/marginal1d/utils.py |  3 +-
 .../ppl/diagnostics/tools/utils/base.py       |  2 +-
 .../ppl/diagnostics/tools/utils/hdi.py        | 90 -------------------
 4 files changed, 12 insertions(+), 93 deletions(-)
 delete mode 100644 src/beanmachine/ppl/diagnostics/tools/utils/hdi.py

diff --git a/src/beanmachine/ppl/diagnostics/tools/js/yarn.lock b/src/beanmachine/ppl/diagnostics/tools/js/yarn.lock
index 5ece43a2c8..439575393b 100644
--- a/src/beanmachine/ppl/diagnostics/tools/js/yarn.lock
+++ b/src/beanmachine/ppl/diagnostics/tools/js/yarn.lock
@@ -836,6 +836,11 @@ emoji-regex@^9.2.2:
   resolved "https://registry.yarnpkg.com/emoji-regex/-/emoji-regex-9.2.2.tgz#840c8803b0d8047f4ff0cf963176b32d4ef3ed72"
   integrity sha512-L18DaJsXSUk2+42pv8mLs5jJT2hqFkFE4j21wOmgbUqsZ2hL72NsUU785g9RXgo3s0ZNgVl42TiHp3ZtOv/Vyg==
 
+emoji-regex@^9.2.2:
+  version "9.2.2"
+  resolved "https://registry.yarnpkg.com/emoji-regex/-/emoji-regex-9.2.2.tgz#840c8803b0d8047f4ff0cf963176b32d4ef3ed72"
+  integrity sha512-L18DaJsXSUk2+42pv8mLs5jJT2hqFkFE4j21wOmgbUqsZ2hL72NsUU785g9RXgo3s0ZNgVl42TiHp3ZtOv/Vyg==
+
 enhanced-resolve@^5.0.0, enhanced-resolve@^5.10.0:
   version "5.10.0"
   resolved "https://registry.yarnpkg.com/enhanced-resolve/-/enhanced-resolve-5.10.0.tgz#0dc579c3bb2a1032e357ac45b8f3a6f3ad4fb1e6"
@@ -1675,6 +1680,11 @@ js-sdsl@^4.1.4:
   resolved "https://registry.yarnpkg.com/js-tokens/-/js-tokens-4.0.0.tgz#19203fb59991df98e3a287050d4647cdeaf32499"
   integrity sha512-RdJUflcE3cUzKiMqQgsCu06FPu9UdIJO0beYbPhHN4k6apgJtifcoCtT9bcxOpYBtpD2kCM6Sbzg4CausW/PKQ==
 
+"js-tokens@^3.0.0 || ^4.0.0":
+  version "4.0.0"
+  resolved "https://registry.yarnpkg.com/js-tokens/-/js-tokens-4.0.0.tgz#19203fb59991df98e3a287050d4647cdeaf32499"
+  integrity sha512-RdJUflcE3cUzKiMqQgsCu06FPu9UdIJO0beYbPhHN4k6apgJtifcoCtT9bcxOpYBtpD2kCM6Sbzg4CausW/PKQ==
+
 js-yaml@^4.1.0:
   version "4.1.0"
   resolved "https://registry.yarnpkg.com/js-yaml/-/js-yaml-4.1.0.tgz#c1fb65f8f5017901cdd2c951864ba18458a10602"
diff --git a/src/beanmachine/ppl/diagnostics/tools/marginal1d/utils.py b/src/beanmachine/ppl/diagnostics/tools/marginal1d/utils.py
index 2061c890ea..6ef4dfebf7 100644
--- a/src/beanmachine/ppl/diagnostics/tools/marginal1d/utils.py
+++ b/src/beanmachine/ppl/diagnostics/tools/marginal1d/utils.py
@@ -8,9 +8,8 @@
 from typing import List
 
 import numpy as np
-import statsmodels.api as sm
 from beanmachine.ppl.diagnostics.tools.marginal1d import typing
-from beanmachine.ppl.diagnostics.tools.utils import hdi, plotting_utils
+from beanmachine.ppl.diagnostics.tools.utils import plotting_utils
 from bokeh.models.annotations import Band, LabelSet
 from bokeh.models.glyphs import Circle, Line
 from bokeh.models.layouts import Column, Row
diff --git a/src/beanmachine/ppl/diagnostics/tools/utils/base.py b/src/beanmachine/ppl/diagnostics/tools/utils/base.py
index 8f290ec068..3a624c219b 100644
--- a/src/beanmachine/ppl/diagnostics/tools/utils/base.py
+++ b/src/beanmachine/ppl/diagnostics/tools/utils/base.py
@@ -23,7 +23,7 @@ class Base(ABC):
     @abstractmethod
     def __init__(self: T, mcs: MonteCarloSamples) -> None:
         self.data = serialize_bm(mcs)
-        self.rv_names = list(self.data.keys())
+        self.rv_names = ["Select a random variable..."] + list(self.data.keys())
         self.num_chains = mcs.num_chains
         self.num_draws = mcs.get_num_samples()
         self.palette = plotting_utils.choose_palette(self.num_chains)
diff --git a/src/beanmachine/ppl/diagnostics/tools/utils/hdi.py b/src/beanmachine/ppl/diagnostics/tools/utils/hdi.py
deleted file mode 100644
index d14594beff..0000000000
--- a/src/beanmachine/ppl/diagnostics/tools/utils/hdi.py
+++ /dev/null
@@ -1,90 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-#
-# This source code is licensed under the MIT license found in the
-# LICENSE file in the root directory of this source tree.
-
-import math
-from typing import Dict, List, Union
-
-
-def interval(
-    data: List[float],
-    hdi_probability: float,
-) -> Dict[str, Union[List[float], float]]:
-    """Find the lower and upper bound of the Highest Density Interval (HDI).
-
-    Parameters
-    ----------
-    data : List[float]
-        Raw random variable data from the model.
-    hdi_probability : float
-        The highest density interval probability to use when calculating the HDI.
-
-    Returns
-    -------
-    Dict[str, Union[List[float], float]]
-        A dictionary defining the lower and upper bound for the HDI of the given data,
-        as well as the indices of the lower and upper bound for the sorted raw data.
-    """
-    N = len(data)
-    sorted_data = sorted(data)
-    stop_index = math.floor(hdi_probability * N)
-    start_index = N - stop_index
-    left_data = sorted_data[stop_index:]
-    right_data = sorted_data[:start_index]
-    hdi = []
-    for i in range(len(left_data)):
-        hdi.append(left_data[i] - right_data[i])
-    lower_index = hdi.index(min(hdi))
-    upper_index = lower_index + stop_index
-    return {
-        "lower_bound": sorted_data[lower_index],
-        "upper_bound": sorted_data[upper_index],
-        "lower_bound_index": lower_index,
-        "upper_bound_index": upper_index,
-    }
-
-
-def data(
-    rv_data: List[float],
-    marginal_x: List[float],
-    marginal_y: List[float],
-    hdi_probability: float,
-) -> Dict[str, Union[List[float], float]]:
-    """Construct x and y arrays from the HDI bounds using the raw random variable dataself.
-
-    It also returns the actual HDI bounds of the random variable data.
-
-    Parameters
-    ----------
-    rv_data : List[float]
-        Raw random variable data from the model.
-    marginalX : List[float]
-        The support of the Kernel Density Estimate of the random variable.
-    marginalY : List[float]
-        The Kernel Density Estimate of the random variable.
-    hdiProbability : float
-    The highest density interval probability to use when alculating the HDI.
-
-    Returns
-    -------
-    Dict[str, Union[List[float], float]]
-        The x array for the HDI data is the "lower" key, while the y array for the HDI
-        data is the "upper" key and the "base" key is the support along the axis the HDI
-        is calculated. This nomenclature mirrors the Bokeh Band object, which is the
-        annotation type we will use when rendering the HDI intervals.
-    """
-    hdi = interval(rv_data, hdi_probability)
-    base = []
-    upper = []
-    for i in range(len(marginal_x)):
-        if marginal_x[i] <= hdi["upper_bound"] and marginal_x[i] >= hdi["lower_bound"]:
-            base.append(marginal_x[i])
-            upper.append(marginal_y[i])
-    return {
-        "base": base,
-        "lower": [0] * len(base),
-        "upper": upper,
-        "lower_bound": hdi["lower_bound"],
-        "upper_bound": hdi["upper_bound"],
-    }

From 1d32e056aacaf6d522c075071e8f5f7bb4b63ed0 Mon Sep 17 00:00:00 2001
From: Andy Maloney <amaloney@quansight.com>
Date: Fri, 30 Sep 2022 11:44:45 -0500
Subject: [PATCH 4/7] Trace diagnostic tool

This tool branches from the marginal1d tool and rebases against main.
Added files include those needed for displaying a Trace diagnostic tool
for Bean Machine models.

This also updates the Coin_flipping.ipynb tutorial to include the new
diagnostic tool.
---
 .../ppl/diagnostics/tools/__init__.py         |   6 +-
 .../diagnostics/tools/js/src/stats/array.ts   |  89 +++
 .../tools/js/src/stats/dataTransformation.ts  |  48 ++
 .../tools/js/src/stats/histogram.ts           | 146 ++++
 .../tools/js/src/trace/callbacks.ts           | 220 ++++++
 .../diagnostics/tools/js/src/trace/index.ts   |  12 +
 .../tools/js/src/trace/interfaces.ts          | 110 +++
 .../ppl/diagnostics/tools/js/tsconfig.json    |   1 -
 .../diagnostics/tools/js/webpack.config.js    |  15 +-
 .../ppl/diagnostics/tools/js/yarn.lock        |  90 +--
 .../ppl/diagnostics/tools/trace/__init__.py   |   4 +
 .../ppl/diagnostics/tools/trace/tool.py       | 168 ++++
 .../ppl/diagnostics/tools/trace/typing.py     | 158 ++++
 .../ppl/diagnostics/tools/trace/utils.py      | 745 ++++++++++++++++++
 .../ppl/diagnostics/tools/utils/base.py       |   6 +-
 src/beanmachine/ppl/diagnostics/tools/viz.py  |  14 +
 16 files changed, 1738 insertions(+), 94 deletions(-)
 create mode 100644 src/beanmachine/ppl/diagnostics/tools/js/src/stats/histogram.ts
 create mode 100644 src/beanmachine/ppl/diagnostics/tools/js/src/trace/callbacks.ts
 create mode 100644 src/beanmachine/ppl/diagnostics/tools/js/src/trace/index.ts
 create mode 100644 src/beanmachine/ppl/diagnostics/tools/js/src/trace/interfaces.ts
 create mode 100644 src/beanmachine/ppl/diagnostics/tools/trace/__init__.py
 create mode 100644 src/beanmachine/ppl/diagnostics/tools/trace/tool.py
 create mode 100644 src/beanmachine/ppl/diagnostics/tools/trace/typing.py
 create mode 100644 src/beanmachine/ppl/diagnostics/tools/trace/utils.py

diff --git a/src/beanmachine/ppl/diagnostics/tools/__init__.py b/src/beanmachine/ppl/diagnostics/tools/__init__.py
index 3e58d695f5..9cd2d21afd 100644
--- a/src/beanmachine/ppl/diagnostics/tools/__init__.py
+++ b/src/beanmachine/ppl/diagnostics/tools/__init__.py
@@ -12,9 +12,13 @@
 
 
 if sys.version_info >= (3, 8):
+    # NOTE: We need to import NotRequired from typing_extensions until PEP 655 is
+    #       accepted, see https://peps.python.org/pep-0655/. This is to follow the
+    #       interface objects in JavaScript that allow keys to not be required using ?.
     from typing import TypedDict
+    from typing_extensions import NotRequired
 else:
-    from typing_extensions import TypedDict
+    from typing_extensions import NotRequired, TypedDict
 
 
 TOOLS_DIR = Path(__file__).parent.resolve()
diff --git a/src/beanmachine/ppl/diagnostics/tools/js/src/stats/array.ts b/src/beanmachine/ppl/diagnostics/tools/js/src/stats/array.ts
index d67b2e7dbf..39dcf65bd0 100644
--- a/src/beanmachine/ppl/diagnostics/tools/js/src/stats/array.ts
+++ b/src/beanmachine/ppl/diagnostics/tools/js/src/stats/array.ts
@@ -39,3 +39,92 @@ export const numericalSort = (data: number[]): number[] => {
     return a < b ? -1 : a > b ? 1 : 0;
   });
 };
+
+/**
+ * Determine the shape of the given array.
+ *
+ * @param {any[]} data - Any array of data.
+ * @returns {number[]} The shape of the data as an array.
+ */
+export const shape = (data: any[]): number[] => {
+  // From https://stackoverflow.com/questions/10237615/get-size-of-dimensions-in-array
+  const computeShape = (array: any[]): any[] => {
+    return array.length ? [...[array.length], ...computeShape(array[0])] : [];
+  };
+  const arrayShape = computeShape(data);
+  // Remove the empty array that will exist at the end of the shape array, since it is
+  // the returned "else" value from above.
+  const dataShape = [];
+  for (let i = 0; i < arrayShape.length; i += 1) {
+    if (!Array.isArray(arrayShape[i])) {
+      dataShape.push(arrayShape[i]);
+    }
+  }
+  return dataShape;
+};
+
+/**
+ * Create an array that starts and stops with the given number of steps.
+ *
+ * @param {number} start - Where to start the array from.
+ * @param {number} stop - Where to stop the array.
+ * @param {number} [step] - The step size to take.
+ * @param {boolean} [closed] - Flag used to return a closed array or not.
+ * @param {null | number} [size] - If not null, then will return an array with the given
+ *     size.
+ * @returns {number[]} An array that is linearly spaced between the start and stop
+ *     values.
+ */
+export const linearRange = (
+  start: number,
+  stop: number,
+  step: number = 1,
+  closed: boolean = true,
+  size: null | number = null,
+): number[] => {
+  if (size !== null) {
+    step = (stop - start) / size;
+  }
+  let len = (stop - start) / step + 1;
+  if (!closed) {
+    len = (stop - start - step) / step + 1;
+  }
+  return Array.from({length: len}, (_, i) => {
+    return start + i * step;
+  });
+};
+
+/**
+ * Return the indices that would sort the array. Follows NumPy's implementation.
+ *
+ * @param {number[]} data - The data to sort.
+ * @returns {number[]} An array of indices that would sort the original array.
+ */
+export const argSort = (data: number[]): number[] => {
+  const dataCopy = data.slice(0);
+  return dataCopy
+    .map((value, index) => {
+      return [value, index];
+    })
+    .sort((a, b) => {
+      return a[0] - b[0];
+    })
+    .map((value) => {
+      return value[1];
+    });
+};
+
+/**
+ * Count the number of time a value appears in an array.
+ *
+ * @param {number[]} data - The numeric array to count objects for.
+ * @returns {{[key: string]: number}} An object that contains the keys as the items in
+ *     the original array, and values that are counts of the key.
+ */
+export const valueCounts = (data: number[]): {[key: string]: number} => {
+  const counts: {[key: string]: number} = {};
+  for (let i = 0; i < data.length; i += 1) {
+    counts[data[i]] = (counts[data[i]] || 0) + 1;
+  }
+  return counts;
+};
diff --git a/src/beanmachine/ppl/diagnostics/tools/js/src/stats/dataTransformation.ts b/src/beanmachine/ppl/diagnostics/tools/js/src/stats/dataTransformation.ts
index a8ecb46596..2c482a1b41 100644
--- a/src/beanmachine/ppl/diagnostics/tools/js/src/stats/dataTransformation.ts
+++ b/src/beanmachine/ppl/diagnostics/tools/js/src/stats/dataTransformation.ts
@@ -5,6 +5,8 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+import {argSort, valueCounts} from './array';
+
 /**
  * Scale the given array of numbers by the given scaleFactor. Note that this method
  * divides values in the given array by the scaleFactor.
@@ -32,3 +34,49 @@ export const scaleToOne = (data: number[]): number[] => {
   const scaleFactor = Math.max(...data);
   return scaleBy(data, scaleFactor);
 };
+
+/**
+ * Assign ranks to the given data. Follows SciPy's and ArviZ's implementations.
+ *
+ * @param {number[]} data - The numeric data to rank.
+ * @returns {number[]} An array of rankings.
+ */
+export const rankData = (data: number[]): number[] => {
+  const n = data.length;
+  const rank = Array(n);
+  const sortedIndex = argSort(data);
+  for (let i = 0; i < rank.length; i += 1) {
+    rank[sortedIndex[i]] = i + 1;
+  }
+  const counts = valueCounts(data);
+  const countsArray = Object.entries(counts);
+  const keys = [];
+  const keyCounts = [];
+  for (let i = 0; i < countsArray.length; i += 1) {
+    const [key, count] = countsArray[i];
+    if (count > 1) {
+      keys.push(parseFloat(key));
+      keyCounts.push(count);
+    }
+  }
+  for (let i = 0; i < keys.length; i += 1) {
+    const repeatIndices = [];
+    for (let j = 0; j < data.length; j += 1) {
+      if (data[j] === keys[i]) {
+        repeatIndices.push(j);
+      }
+    }
+    const rankValues = [];
+    for (let k = 0; k < repeatIndices.length; k += 1) {
+      rankValues.push(rank[repeatIndices[k]]);
+    }
+    const sum = rankValues.reduce((previousValue, currentValue) => {
+      return previousValue + currentValue;
+    }, 0.0);
+    const rankMean = sum / rankValues.length;
+    for (let k = 0; k < repeatIndices.length; k += 1) {
+      rank[repeatIndices[k]] = rankMean;
+    }
+  }
+  return rank;
+};
diff --git a/src/beanmachine/ppl/diagnostics/tools/js/src/stats/histogram.ts b/src/beanmachine/ppl/diagnostics/tools/js/src/stats/histogram.ts
new file mode 100644
index 0000000000..394a632c24
--- /dev/null
+++ b/src/beanmachine/ppl/diagnostics/tools/js/src/stats/histogram.ts
@@ -0,0 +1,146 @@
+/* import {calculateHistogram} from 'compute-histogram'; */
+import {linearRange, numericalSort, shape} from './array';
+import {rankData, scaleToOne} from './dataTransformation';
+import {mean as computeMean} from './pointStatistic';
+
+/**
+ * Compute the histogram of the given data.
+ *
+ * @param {number[]} data - Data to bin.
+ * @param {number} [numBins] - The number of bins to use for the histogram. If none is
+ *     given, then we follow ArviZ's implementation by using twice then number of bins
+ *     of the Sturges formula.
+ * @returns {number[][]} [TODO:description]
+ */
+export const calculateHistogram = (data: number[], numBins: number = 0): number[][] => {
+  const sortedData = numericalSort(data);
+  const numSamples = sortedData.length;
+  const dataMin = Math.min(...data);
+  const dataMax = Math.max(...data);
+  if (numBins === 0) {
+    numBins = Math.floor(Math.ceil(2 * Math.log2(numSamples)) + 1);
+  }
+  const binSize =
+    (dataMax - dataMin) / numBins === 0 ? 1 : (dataMax - dataMin) / numBins;
+  const bins = Array(numBins)
+    .fill([0, 0])
+    .map((_, i) => {
+      return [i, 0];
+    });
+
+  for (let i = 0; i < data.length; i += 1) {
+    const datum = sortedData[i];
+    let binIndex = Math.floor((datum - dataMin) / binSize);
+    // Subtract 1 if the value lies on the last bin.
+    if (binIndex === numBins) {
+      binIndex -= 1;
+    }
+    bins[binIndex][1] += 1;
+  }
+  return bins;
+};
+
+export interface RankHistogram {
+  [key: string]: {
+    quad: {
+      left: number[];
+      top: number[];
+      right: number[];
+      bottom: number[];
+      chain: number[];
+      draws: string[];
+      rank: number[];
+    };
+    line: {x: number[]; y: number[]};
+    chain: number[];
+    rankMean: number[];
+    mean: number[];
+  };
+}
+
+/**
+ * A histogram of rank data.
+ *
+ * @param {number[][]} data - Raw random variable data for several chains.
+ * @returns {RankHistogram} A histogram of the data rankings.
+ */
+export const rankHistogram = (data: number[][]): RankHistogram => {
+  const [numChains, numDraws] = shape(data);
+  const numSamples = numChains * numDraws;
+  const flatData = data.flat();
+
+  // Calculate the rank of the data and ensure it is the same shape as the original
+  // data.
+  const rank = rankData(flatData);
+  const rankArray = [];
+  let start = Number.NaN;
+  let end = Number.NaN;
+  for (let i = 0; i < numChains; i += 1) {
+    if (i === 0) {
+      start = 0;
+      end = numDraws;
+    } else {
+      start = end;
+      end = (i + 1) * numDraws;
+    }
+    const chainRanks = rank.slice(start, end);
+    rankArray.push(chainRanks);
+    start = end;
+    end = (i + 1) * numDraws;
+  }
+
+  // Calculate the number of bins needed. We will follow ArviZ and use twice the result
+  // using the Sturges' formula.
+  const numBins = Math.floor(Math.ceil(2 * Math.log2(numSamples)) + 1);
+  const lastBinEdge = Math.max(...rank);
+
+  // Calculate the bin edges. Since the linearRange function computes a linear spacing
+  // of values between the start and end point, we need to ensure they are integer
+  // values.
+  let binEdges = linearRange(0, lastBinEdge, 1, true, numBins);
+  binEdges = binEdges.map((value) => {
+    return Math.ceil(value);
+  });
+
+  // Calculate the histograms of the rank data, and normalize it for each chain.
+  const output = {} as RankHistogram;
+  for (let i = 0; i < numChains; i += 1) {
+    const chainIndex = i + 1;
+    const chainName = `chain${chainIndex}`;
+    const chainRankHistogram = calculateHistogram(rankArray[i], numBins);
+    let counts = [];
+    for (let j = 0; j < chainRankHistogram.length; j += 1) {
+      counts.push(chainRankHistogram[j][1]);
+    }
+    counts = scaleToOne(counts);
+    const chainCounts = counts.map((value) => {
+      return value + i;
+    });
+
+    const chainRankMean = computeMean(chainCounts);
+    const left = binEdges.slice(0, binEdges.length - 1);
+    const right = binEdges.slice(1);
+    const binLabel = [];
+    for (let j = 0; j < left.length; j += 1) {
+      binLabel.push(`${left[j].toLocaleString()}-${right[j].toLocaleString()}`);
+    }
+    const x = linearRange(0, numSamples, 1);
+    const y = Array(x.length).fill(chainRankMean);
+    output[chainName] = {
+      quad: {
+        left: left,
+        top: chainCounts,
+        right: right,
+        bottom: Array(numBins).fill(i),
+        chain: Array(left.length).fill(i + 1),
+        draws: binLabel,
+        rank: counts,
+      },
+      line: {x: x, y: y},
+      chain: Array(x.length).fill(i + 1),
+      rankMean: Array(x.length).fill(chainIndex - chainRankMean),
+      mean: Array(x.length).fill(computeMean(counts)),
+    };
+  }
+  return output;
+};
diff --git a/src/beanmachine/ppl/diagnostics/tools/js/src/trace/callbacks.ts b/src/beanmachine/ppl/diagnostics/tools/js/src/trace/callbacks.ts
new file mode 100644
index 0000000000..1c4d4ab362
--- /dev/null
+++ b/src/beanmachine/ppl/diagnostics/tools/js/src/trace/callbacks.ts
@@ -0,0 +1,220 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import {Axis} from '@bokehjs/models/axes/axis';
+import * as interfaces from './interfaces';
+import {linearRange, shape} from '../stats/array';
+import {interval as hdiInterval} from '../stats/highestDensityInterval';
+import {rankHistogram} from '../stats/histogram';
+import {oneD} from '../stats/marginal';
+import {mean} from '../stats/pointStatistic';
+
+const figureNames = ['marginals', 'forests', 'traces', 'ranks'];
+
+/**
+ * Update the given Bokeh Axis object with the new label string. You must use this
+ * method to update axis strings using TypeScript, otherwise the ts compiler will throw
+ * a type check error.
+ *
+ * @param {Axis} axis - The Bokeh Axis object needing a new label.
+ * @param {string | null} label - The new label for the Bokeh Axis object.
+ */
+export const updateAxisLabel = (axis: Axis, label: string | null): void => {
+  // Type check requirement.
+  if ('axis_label' in axis) {
+    axis.axis_label = label;
+  }
+};
+
+/**
+ * Compute data for the trace diagnostic tool.
+ *
+ * @param {number[][]} data - Raw random variable data from the model for all chains.
+ * @param {number} bwFactor - Multiplicative factor to be applied to the bandwidth when
+ *     calculating the Kernel Density Estimate (KDE).
+ * @param {number} hdiProbability - The highest density interval probability to use when
+ *     calculating the HDI.
+ * @returns {interfaces.Data} Data object that contains data for each figure including
+ *     each chain.
+ */
+export const computeData = (
+  data: number[][],
+  bwFactor: number,
+  hdiProbability: number,
+): interfaces.Data => {
+  const [numChains, numDraws] = shape(data);
+  const output = {} as interfaces.Data;
+  for (let i = 0; i < figureNames.length; i += 1) {
+    const figureName = figureNames[i];
+    if (figureName !== 'ranks') {
+      switch (figureName) {
+        case 'marginals':
+          output[figureName] = {} as interfaces.MarginalDataAllChains;
+          break;
+        case 'forests':
+          output[figureName] = {} as interfaces.ForestDataAllChains;
+          break;
+        case 'traces':
+          output[figureName] = {} as interfaces.TraceDataAllChains;
+          break;
+        default:
+          break;
+      }
+      for (let j = 0; j < numChains; j += 1) {
+        const chainIndex = j + 1;
+        const chainName = `chain${chainIndex}`;
+        const chainData = data[j];
+        const marginal = oneD(chainData, bwFactor);
+        const marginalMean = mean(marginal.x);
+        let hdiBounds;
+        switch (figureName) {
+          case 'marginals':
+            output[figureName][chainName] = {} as interfaces.MarginalDataSingleChain;
+            output[figureName][chainName] = {
+              line: {x: marginal.x, y: marginal.y},
+              chain: chainIndex,
+              mean: marginalMean,
+              bandwidth: marginal.bandwidth,
+            };
+            break;
+          case 'forests':
+            output[figureName][chainName] = {} as interfaces.ForestDataSingleChain;
+            hdiBounds = hdiInterval(chainData, hdiProbability);
+            output[figureName][chainName] = {
+              line: {
+                x: [hdiBounds.lowerBound, hdiBounds.upperBound],
+                y: Array(2).fill(chainIndex),
+              },
+              circle: {x: [marginalMean], y: [chainIndex]},
+              chain: chainIndex,
+              mean: marginalMean,
+            };
+            break;
+          case 'traces':
+            output[figureName][chainName] = {} as interfaces.TraceDataSingleChain;
+            output[figureName][chainName] = {
+              line: {x: linearRange(0, numDraws - 1, 1), y: chainData},
+              chain: chainIndex,
+              mean: marginalMean,
+            };
+            break;
+          default:
+            break;
+        }
+      }
+    } else if (figureName === 'ranks') {
+      output[figureName] = rankHistogram(data);
+    }
+  }
+  return output;
+};
+
+/**
+ * Callback used to update the Bokeh application in the notebook.
+ *
+ * @param {number[][]} data - Raw random variable data from the model for all chains.
+ * @param {string} rvName - The name of the random variable from the model.
+ * @param {number} bwFactor - Multiplicative factor to be applied to the bandwidth when
+ *     calculating the kernel density estimate.
+ * @param {number} hdiProbability - The highest density interval probability to use when
+ *     calculating the HDI.
+ * @param {interfaces.Sources} sources - Bokeh sources used to render glyphs in the
+ *     application.
+ * @param {interfaces.Figures} figures - Bokeh figures shown in the application.
+ * @param {interfaces.Tooltips} tooltips - Bokeh tooltips shown on the glyphs.
+ */
+export const update = (
+  data: number[][],
+  rvName: string,
+  bwFactor: number,
+  hdiProbability: number,
+  sources: interfaces.Sources,
+  figures: interfaces.Figures,
+  tooltips: interfaces.Tooltips,
+): void => {
+  const [numChains] = shape(data);
+  const computedData = computeData(data, bwFactor, hdiProbability);
+  for (let i = 0; i < figureNames.length; i += 1) {
+    const figureName = figureNames[i];
+    const figure = figures[figureName];
+    for (let j = 0; j < numChains; j += 1) {
+      const chainIndex = j + 1;
+      const chainName = `chain${chainIndex}`;
+      const chainData = computedData[figureName][chainName];
+      const source = sources[figureName][chainName];
+      switch (figureName) {
+        case 'marginals':
+          source.line.data = {
+            x: chainData.line.x,
+            y: chainData.line.y,
+            chain: Array(chainData.line.x.length).fill(chainData.chain),
+            mean: Array(chainData.line.x.length).fill(chainData.mean),
+          };
+          updateAxisLabel(figure.below[0], rvName);
+          tooltips[figureName][j].tooltips = [
+            ['Chain', '@chain'],
+            ['Mean', '@mean'],
+            [rvName, '@x'],
+          ];
+          break;
+        case 'forests':
+          source.line.data = {
+            x: chainData.line.x,
+            y: chainData.line.y,
+            chain: Array(chainData.line.x.length).fill(chainData.chain),
+            mean: Array(chainData.line.x.length).fill(chainData.mean),
+          };
+          source.circle.data = {
+            x: chainData.circle.x,
+            y: chainData.circle.y,
+            chain: [chainData.chain],
+            mean: [chainData.mean],
+          };
+          updateAxisLabel(figure.below[0], rvName);
+          tooltips[figureName][j].tooltips = [
+            ['Chain', '@chain'],
+            [rvName, '@mean'],
+          ];
+          break;
+        case 'traces':
+          source.line.data = {
+            x: chainData.line.x,
+            y: chainData.line.y,
+            chain: Array(chainData.line.x.length).fill(chainData.chain),
+            mean: Array(chainData.line.x.length).fill(chainData.mean),
+          };
+          updateAxisLabel(figure.left[0], rvName);
+          tooltips[figureName][j].tooltips = [
+            ['Chain', '@chain'],
+            ['Mean', '@mean'],
+            [rvName, '@y'],
+          ];
+          break;
+        case 'ranks':
+          source.line.data = {
+            x: chainData.line.x,
+            y: chainData.line.y,
+            chain: chainData.chain,
+            rankMean: chainData.rankMean,
+          };
+          tooltips[figureName][j].line.tooltips = [
+            ['Chain', '@chain'],
+            ['Rank mean', '@rankMean'],
+          ];
+          source.quad.data = chainData.quad;
+          tooltips[figureName][j].quad.tooltips = [
+            ['Chain', '@chain'],
+            ['Draws', '@draws'],
+            ['Rank', '@rank'],
+          ];
+          break;
+        default:
+          break;
+      }
+    }
+  }
+};
diff --git a/src/beanmachine/ppl/diagnostics/tools/js/src/trace/index.ts b/src/beanmachine/ppl/diagnostics/tools/js/src/trace/index.ts
new file mode 100644
index 0000000000..8b75182970
--- /dev/null
+++ b/src/beanmachine/ppl/diagnostics/tools/js/src/trace/index.ts
@@ -0,0 +1,12 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import * as trace from './callbacks';
+
+// The CustomJS methods used by Bokeh require us to make the JavaScript available in the
+// browser, which is done by defining it below.
+(window as any).trace = trace;
diff --git a/src/beanmachine/ppl/diagnostics/tools/js/src/trace/interfaces.ts b/src/beanmachine/ppl/diagnostics/tools/js/src/trace/interfaces.ts
new file mode 100644
index 0000000000..b7e1961889
--- /dev/null
+++ b/src/beanmachine/ppl/diagnostics/tools/js/src/trace/interfaces.ts
@@ -0,0 +1,110 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import {ColumnDataSource} from '@bokehjs/models/sources/column_data_source';
+import {HoverTool} from '@bokehjs/models/tools/inspectors/hover_tool';
+import {Plot} from '@bokehjs/models/plots/plot';
+import {RankHistogram} from '../stats/histogram';
+
+// NOTE: In the corresponding Python typing files for the diagnostic tool, we define
+//       similar types using a TypedDict object. TypeScript allows us to maintain
+//       semantic information about the key names and their types in the same object and
+//       still have the ability to loop over objects as long as we have the
+//       [key: string]: any; indicator in the interface definition. This boils down to a
+//       Python type of Dict[Any, Any], which again loses all type information about the
+//       object we are defining. We are mirroring what is done in Python here, so we
+//       keep the semantic information here at the expense of losing type information
+//       similarly to what is done in Python.
+
+export interface LineOrCircleGlyphData {
+  [key: string]: any;
+  x: number[];
+  y: number[];
+}
+
+export interface MarginalDataSingleChain {
+  [key: string]: any;
+  line: LineOrCircleGlyphData;
+  chain: number;
+  mean: number;
+  bandwidth: number;
+}
+
+export interface ForestDataSingleChain {
+  [key: string]: any;
+  line: LineOrCircleGlyphData;
+  circle: LineOrCircleGlyphData;
+  chain: number;
+  mean: number;
+}
+
+export interface TraceDataSingleChain {
+  [key: string]: any;
+  line: LineOrCircleGlyphData;
+  chain: number;
+  mean: number;
+}
+
+export interface MarginalDataAllChains {
+  [key: string]: MarginalDataSingleChain;
+}
+
+export interface ForestDataAllChains {
+  [key: string]: ForestDataSingleChain;
+}
+
+export interface TraceDataAllChains {
+  [key: string]: TraceDataSingleChain;
+}
+
+export interface Data {
+  [key: string]: any;
+  marginals: MarginalDataAllChains;
+  forests: ForestDataAllChains;
+  traces: TraceDataAllChains;
+  ranks: RankHistogram;
+}
+
+export interface SourceSingleChain {
+  line: ColumnDataSource;
+  circle?: ColumnDataSource;
+  quad?: ColumnDataSource;
+}
+
+export interface SourceAllChains {
+  [key: string]: SourceSingleChain;
+}
+
+export interface Sources {
+  [key: string]: any;
+  marginals: SourceAllChains;
+  forests: SourceAllChains;
+  traces: SourceAllChains;
+  ranks: SourceAllChains;
+}
+
+export interface Figures {
+  [key: string]: any;
+  marginals: Plot;
+  forests: Plot;
+  traces: Plot;
+  ranks: Plot;
+}
+
+export interface RankTooltips {
+  [key: string]: HoverTool;
+  line: HoverTool;
+  quad: HoverTool;
+}
+
+export interface Tooltips {
+  [key: string]: any;
+  marginals: HoverTool[];
+  forests: HoverTool[];
+  traces: HoverTool[];
+  ranks: RankTooltips[];
+}
diff --git a/src/beanmachine/ppl/diagnostics/tools/js/tsconfig.json b/src/beanmachine/ppl/diagnostics/tools/js/tsconfig.json
index 37de8788f8..2212d273f3 100644
--- a/src/beanmachine/ppl/diagnostics/tools/js/tsconfig.json
+++ b/src/beanmachine/ppl/diagnostics/tools/js/tsconfig.json
@@ -23,7 +23,6 @@
         "node_modules/@bokeh/bokehjs/build/js/lib/*",
         "node_modules/@bokeh/bokehjs/build/js/types/*"
       ],
-      "compute-histogram/*": ["node_modules/compute-histogram/*"],
       "fast-kde/*": ["node_modules/fast-kde/*"],
       "ndarray/*": ["node_modules/ndarray/*"],
       "ndarray-fft/*": [
diff --git a/src/beanmachine/ppl/diagnostics/tools/js/webpack.config.js b/src/beanmachine/ppl/diagnostics/tools/js/webpack.config.js
index cbe79da94d..e13c2507f6 100644
--- a/src/beanmachine/ppl/diagnostics/tools/js/webpack.config.js
+++ b/src/beanmachine/ppl/diagnostics/tools/js/webpack.config.js
@@ -10,24 +10,35 @@ const path = require('path');
 module.exports = {
   entry: {
     marginal1d: './src/marginal1d/index.ts',
+    trace: './src/trace/index.ts',
   },
   output: {
     filename: '[name].js',
     path: path.resolve(__dirname, 'dist'),
   },
   module: {
-    rules: [{test: /\.ts$/, use: 'ts-loader', exclude: /node_modules/}],
+    rules: [
+      {
+        test: /\.ts$/,
+        use: 'ts-loader',
+        exclude: /node_modules\/(?!(@bokeh\/bokehjs\/build\/js\/lib)\/).*/,
+      },
+    ],
   },
   target: 'web',
   mode: 'production',
   resolve: {
-    extensions: ['.ts'],
+    extensions: ['.ts', '.js'],
     modules: ['./stats', './interfaces', './types', './node_modules'],
     alias: {
       'fast-kde/src/density1d': path.resolve(
         __dirname,
         'node_modules/fast-kde/src/density1d.js',
       ),
+      '@bokehjs/models/ranges/range1d': path.resolve(
+        __dirname,
+        'node_modules/@bokeh/bokehjs/build/js/lib/models/ranges/range1d.js',
+      ),
     },
   },
   optimization: {
diff --git a/src/beanmachine/ppl/diagnostics/tools/js/yarn.lock b/src/beanmachine/ppl/diagnostics/tools/js/yarn.lock
index 439575393b..a6ef03786a 100644
--- a/src/beanmachine/ppl/diagnostics/tools/js/yarn.lock
+++ b/src/beanmachine/ppl/diagnostics/tools/js/yarn.lock
@@ -58,13 +58,6 @@
     jquery-ui ">=1.8.0"
     tslib "^1.10.0"
 
-"@cspotcode/source-map-support@^0.8.0":
-  version "0.8.1"
-  resolved "https://registry.yarnpkg.com/@cspotcode/source-map-support/-/source-map-support-0.8.1.tgz#00629c35a688e05a88b1cda684fb9d5e73f000a1"
-  integrity sha512-IchNf6dN4tHoMFIn/7OE8LWZ19Y6q/67Bmf6vnGREv8RSbBVb9LPJxEcnwrcwX6ixSvaiGoomAUvu4YSxXrVgw==
-  dependencies:
-    "@jridgewell/trace-mapping" "0.3.9"
-
 "@discoveryjs/json-ext@^0.5.0":
   version "0.5.7"
   resolved "https://registry.yarnpkg.com/@discoveryjs/json-ext/-/json-ext-0.5.7.tgz#1d572bfbbe14b7704e0ba0f39b74815b84870d70"
@@ -136,14 +129,6 @@
   resolved "https://registry.yarnpkg.com/@jridgewell/sourcemap-codec/-/sourcemap-codec-1.4.14.tgz#add4c98d341472a289190b424efbdb096991bb24"
   integrity sha512-XPSJHWmi394fuUuzDnGz1wiKqWfo1yXecHQMRf2l6hztTO+nPru658AyDngaBe7isIxEkRsPR3FZh+s7iVa4Uw==
 
-"@jridgewell/trace-mapping@0.3.9":
-  version "0.3.9"
-  resolved "https://registry.yarnpkg.com/@jridgewell/trace-mapping/-/trace-mapping-0.3.9.tgz#6534fd5933a53ba7cbf3a17615e273a0d1273ff9"
-  integrity sha512-3Belt6tdc8bPgAtbcmdtNJlirVoTmEb5e2gC94PnkwEW9jI6CAHUeoG85tjWP5WquqfavoMtMwiG4P926ZKKuQ==
-  dependencies:
-    "@jridgewell/resolve-uri" "^3.0.3"
-    "@jridgewell/sourcemap-codec" "^1.4.10"
-
 "@jridgewell/trace-mapping@^0.3.14", "@jridgewell/trace-mapping@^0.3.9":
   version "0.3.16"
   resolved "https://registry.yarnpkg.com/@jridgewell/trace-mapping/-/trace-mapping-0.3.16.tgz#a7982f16c18cae02be36274365433e5b49d7b23f"
@@ -173,26 +158,6 @@
     "@nodelib/fs.scandir" "2.1.5"
     fastq "^1.6.0"
 
-"@tsconfig/node10@^1.0.7":
-  version "1.0.9"
-  resolved "https://registry.yarnpkg.com/@tsconfig/node10/-/node10-1.0.9.tgz#df4907fc07a886922637b15e02d4cebc4c0021b2"
-  integrity sha512-jNsYVVxU8v5g43Erja32laIDHXeoNvFEpX33OK4d6hljo3jDhCBDhx5dhCCTMWUojscpAagGiRkBKxpdl9fxqA==
-
-"@tsconfig/node12@^1.0.7":
-  version "1.0.11"
-  resolved "https://registry.yarnpkg.com/@tsconfig/node12/-/node12-1.0.11.tgz#ee3def1f27d9ed66dac6e46a295cffb0152e058d"
-  integrity sha512-cqefuRsh12pWyGsIoBKJA9luFu3mRxCA+ORZvA4ktLSzIuCUtWVxGIuXigEwO5/ywWFMZ2QEGKWvkZG1zDMTag==
-
-"@tsconfig/node14@^1.0.0":
-  version "1.0.3"
-  resolved "https://registry.yarnpkg.com/@tsconfig/node14/-/node14-1.0.3.tgz#e4386316284f00b98435bf40f72f75a09dabf6c1"
-  integrity sha512-ysT8mhdixWK6Hw3i1V2AeRqZ5WfXg1G43mqoYlM2nc6388Fq5jcXyr5mRsqViLx/GJYdoL0bfXD8nmF+Zn/Iow==
-
-"@tsconfig/node16@^1.0.2":
-  version "1.0.3"
-  resolved "https://registry.yarnpkg.com/@tsconfig/node16/-/node16-1.0.3.tgz#472eaab5f15c1ffdd7f8628bd4c4f753995ec79e"
-  integrity sha512-yOlFc+7UtL/89t2ZhjPvvB/DeAr3r+Dq58IgzsFkOAvVC6NMJXmCGjbptdXdR9qsX7pKcTL+s87FtYREi2dEEQ==
-
 "@types/eslint-scope@^3.7.3":
   version "3.7.4"
   resolved "https://registry.yarnpkg.com/@types/eslint-scope/-/eslint-scope-3.7.4.tgz#37fc1223f0786c39627068a12e94d6e6fc61de16"
@@ -226,7 +191,7 @@
   dependencies:
     "@types/sizzle" "*"
 
-"@types/json-schema@*", "@types/json-schema@^7.0.8", "@types/json-schema@^7.0.9":
+"@types/json-schema@*", "@types/json-schema@^7.0.8":
   version "7.0.11"
   resolved "https://registry.yarnpkg.com/@types/json-schema/-/json-schema-7.0.11.tgz#d421b6c527a3037f7c84433fd2c4229e016863d3"
   integrity sha512-wOuvG1SN4Us4rez+tylwwwCV1psiNVOkJeM3AUWUNWg/jDQY2+HE/444y5gc+jBmRqASOm2Oeh5c1axHobwRKQ==
@@ -492,12 +457,7 @@ acorn-jsx@^5.3.2:
   resolved "https://registry.yarnpkg.com/acorn-jsx/-/acorn-jsx-5.3.2.tgz#7ed5bb55908b3b2f1bc55c6af1653bada7f07937"
   integrity sha512-rq9s+JNhf0IChjtDXxllJ7g41oZk5SlXtp0LHwyA5cejwn7vKmKp4pPri6YEePv2PU65sAsegbXtIinmDFDXgQ==
 
-acorn-walk@^8.1.1:
-  version "8.2.0"
-  resolved "https://registry.yarnpkg.com/acorn-walk/-/acorn-walk-8.2.0.tgz#741210f2e2426454508853a2f44d0ab83b7f69c1"
-  integrity sha512-k+iyHEuPgSw6SbuDpGQM+06HQUa04DZ3o+F6CSzXMvvI5KMvnaEqXe+YVe555R9nn6GPt404fos4wcgpw12SDA==
-
-acorn@^8.4.1, acorn@^8.5.0, acorn@^8.7.1, acorn@^8.8.0:
+acorn@^8.5.0, acorn@^8.7.1, acorn@^8.8.0:
   version "8.8.0"
   resolved "https://registry.yarnpkg.com/acorn/-/acorn-8.8.0.tgz#88c0187620435c7f6015803f5539dae05a9dbea8"
   integrity sha512-QOxyigPVrpZ2GXT+PFyZTl6TtOFc5egxHIP9IlQ+RbupQuX4RkT/Bee4/kQuC02Xkzg84JcT7oLYtDIQxp+v7w==
@@ -529,11 +489,6 @@ ansi-styles@^4.1.0:
   dependencies:
     color-convert "^2.0.1"
 
-arg@^4.1.0:
-  version "4.1.3"
-  resolved "https://registry.yarnpkg.com/arg/-/arg-4.1.3.tgz#269fc7ad5b8e42cb63c896d5666017261c144089"
-  integrity sha512-58S9QDqG0Xx27YwPSt9fJxivjYl432YCwfDMfZ+71RAqUrZef7LrKQZ3LHLOwCS4FLNBplP533Zx895SeOCHvA==
-
 argparse@^2.0.1:
   version "2.0.1"
   resolved "https://registry.yarnpkg.com/argparse/-/argparse-2.0.1.tgz#246f50f3ca78a3240f6c997e8a9bd1eac49e4b38"
@@ -800,11 +755,6 @@ define-properties@^1.1.3, define-properties@^1.1.4:
     has-property-descriptors "^1.0.0"
     object-keys "^1.1.1"
 
-diff@^4.0.1:
-  version "4.0.2"
-  resolved "https://registry.yarnpkg.com/diff/-/diff-4.0.2.tgz#60f3aecb89d5fae520c11aa19efc2bb982aade7d"
-  integrity sha512-58lmxKSA4BNyLz+HHMUzlOEpg09FV+ev6ZMe3vJihgdxzgcwZ8VoEEPmALCZG9LmqfVoNMMKpttIYTVG6uDY7A==
-
 dir-glob@^3.0.1:
   version "3.0.1"
   resolved "https://registry.yarnpkg.com/dir-glob/-/dir-glob-3.0.1.tgz#56dbf73d992a4a93ba1584f4534063fd2e41717f"
@@ -1087,7 +1037,7 @@ eslint-plugin-react@^7.28.0:
     semver "^6.3.0"
     string.prototype.matchall "^4.0.7"
 
-eslint-scope@5.1.1, eslint-scope@^5.1.1:
+eslint-scope@5.1.1:
   version "5.1.1"
   resolved "https://registry.yarnpkg.com/eslint-scope/-/eslint-scope-5.1.1.tgz#e786e59a66cb92b3f6c1fb0d508aab174848f48c"
   integrity sha512-2NxwbF/hZ0KpepYN0cNbo+FN6XoK7GaHlQhgx/hIZl6Va0bF45RQOOwhLIy8lQDbuCiadSLCBnH2CFYquit5bw==
@@ -1785,11 +1735,6 @@ lru-cache@^6.0.0:
   dependencies:
     yallist "^4.0.0"
 
-make-error@^1.1.1:
-  version "1.3.6"
-  resolved "https://registry.yarnpkg.com/make-error/-/make-error-1.3.6.tgz#2eb2e37ea9b67c4891f684a1394799af484cf7a2"
-  integrity sha512-s8UhlNe7vPKomQhC1qFelMokr/Sc3AgNbso3n74mVPA5LTZwkB9NlXf4XPamLxJE8h0gh73rM94xvwRT2CVInw==
-
 mathjax-full@^3.2.0:
   version "3.2.2"
   resolved "https://registry.yarnpkg.com/mathjax-full/-/mathjax-full-3.2.2.tgz#43f02e55219db393030985d2b6537ceae82f1fa7"
@@ -2422,25 +2367,6 @@ ts-loader@^9.3.1:
     micromatch "^4.0.0"
     semver "^7.3.4"
 
-ts-node@^10.9.1:
-  version "10.9.1"
-  resolved "https://registry.yarnpkg.com/ts-node/-/ts-node-10.9.1.tgz#e73de9102958af9e1f0b168a6ff320e25adcff4b"
-  integrity sha512-NtVysVPkxxrwFGUUxGYhfux8k78pQB3JqYBXlLRZgdGUqTO5wU/UyHop5p70iEbGhB7q5KmiZiU0Y3KlJrScEw==
-  dependencies:
-    "@cspotcode/source-map-support" "^0.8.0"
-    "@tsconfig/node10" "^1.0.7"
-    "@tsconfig/node12" "^1.0.7"
-    "@tsconfig/node14" "^1.0.0"
-    "@tsconfig/node16" "^1.0.2"
-    acorn "^8.4.1"
-    acorn-walk "^8.1.1"
-    arg "^4.1.0"
-    create-require "^1.1.0"
-    diff "^4.0.1"
-    make-error "^1.1.1"
-    v8-compile-cache-lib "^3.0.1"
-    yn "3.1.1"
-
 tsconfig-paths@^3.14.1:
   version "3.14.1"
   resolved "https://registry.yarnpkg.com/tsconfig-paths/-/tsconfig-paths-3.14.1.tgz#ba0734599e8ea36c862798e920bcf163277b137a"
@@ -2525,11 +2451,6 @@ uri-js@^4.2.2:
   dependencies:
     punycode "^2.1.0"
 
-v8-compile-cache-lib@^3.0.1:
-  version "3.0.1"
-  resolved "https://registry.yarnpkg.com/v8-compile-cache-lib/-/v8-compile-cache-lib-3.0.1.tgz#6336e8d71965cb3d35a1bbb7868445a7c05264bf"
-  integrity sha512-wa7YjyUGfNZngI/vtK0UHAN+lgDCxBPCylVXGp0zu59Fz5aiGtNXaq3DhIov063MorB+VfufLh3JlF2KdTK3xg==
-
 watchpack@^2.4.0:
   version "2.4.0"
   resolved "https://registry.yarnpkg.com/watchpack/-/watchpack-2.4.0.tgz#fa33032374962c78113f93c7f2fb4c54c9862a5d"
@@ -2652,11 +2573,6 @@ yallist@^4.0.0:
   resolved "https://registry.yarnpkg.com/yallist/-/yallist-4.0.0.tgz#9bb92790d9c0effec63be73519e11a35019a3a72"
   integrity sha512-3wdGidZyq5PB084XLES5TpOSRA3wjXAlIWMhum2kRcv/41Sn2emQ0dycQW4uZXLejwKvg6EsvbdlVL+FYEct7A==
 
-yn@3.1.1:
-  version "3.1.1"
-  resolved "https://registry.yarnpkg.com/yn/-/yn-3.1.1.tgz#1e87401a09d767c1d5eab26a6e4c185182d2eb50"
-  integrity sha512-Ux4ygGWsu2c7isFWe8Yu1YluJmqVhxqK2cLXNQA5AcC3QfbGNpM7fu0Y8b/z16pXLnFxZYvWhd3fhBY9DLmC6Q==
-
 yocto-queue@^0.1.0:
   version "0.1.0"
   resolved "https://registry.yarnpkg.com/yocto-queue/-/yocto-queue-0.1.0.tgz#0294eb3dee05028d31ee1a5fa2c556a6aaf10a1b"
diff --git a/src/beanmachine/ppl/diagnostics/tools/trace/__init__.py b/src/beanmachine/ppl/diagnostics/tools/trace/__init__.py
new file mode 100644
index 0000000000..7bec24cb17
--- /dev/null
+++ b/src/beanmachine/ppl/diagnostics/tools/trace/__init__.py
@@ -0,0 +1,4 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
diff --git a/src/beanmachine/ppl/diagnostics/tools/trace/tool.py b/src/beanmachine/ppl/diagnostics/tools/trace/tool.py
new file mode 100644
index 0000000000..3460c92bd8
--- /dev/null
+++ b/src/beanmachine/ppl/diagnostics/tools/trace/tool.py
@@ -0,0 +1,168 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Trace diagnostic tool for a Bean Machine model."""
+
+from typing import Optional, TypeVar
+
+from beanmachine.ppl.diagnostics.tools.trace import utils
+from beanmachine.ppl.diagnostics.tools.utils.base import Base
+from beanmachine.ppl.inference.monte_carlo_samples import MonteCarloSamples
+from bokeh.embed import file_html
+from bokeh.models.callbacks import CustomJS
+from bokeh.resources import INLINE
+
+
+T = TypeVar("T", bound="Trace")
+
+
+class Trace(Base):
+    """Trace tool.
+
+    Parameters
+    ----------
+    mcs : MonteCarloSamples
+        Bean Machine model object.
+
+    Attributes
+    ----------
+    data : Dict[str, List[List[float]]]
+        JSON serialized version of the Bean Machine model.
+    rv_names : List[str]
+        The list of random variables string names for the given model.
+    num_chains : int
+        The number of chains of the model.
+    num_draws : int
+        The number of draws of the model for each chain.
+    palette : List[str]
+        A list of color values used for the glyphs in the figures. The colors are
+        specifically chosen from the Colorblind palette defined in Bokeh.
+    js : str
+        The JavaScript callbacks needed to render the Bokeh tool independently from
+        a Python server.
+    name : str
+        The name to use when saving Bokeh JSON to disk.
+    """
+
+    def __init__(self: T, mcs: MonteCarloSamples) -> None:
+        super(Trace, self).__init__(mcs)
+
+    def create_document(self: T, name: Optional[str] = None) -> str:
+        if name is not None:
+            self.name = name
+
+        # Initialize widget values using Python.
+        rv_name = self.rv_names[0]
+
+        # NOTE: We are going to use Python and Bokeh to render the tool in the notebook
+        #       output cell, however, we WILL NOT use Python to calculate any of the
+        #       statistics displayed in the tool. We do this so we can make the BROWSER
+        #       run all the calculations based on user interactions. If we did not
+        #       employ this strategy, then the initial display a user would receive
+        #       would be calculated by Python, and any subsequent updates would be
+        #       calculated by JavaScript. The side-effect of having two backends
+        #       calculate data could cause the figures to flicker, which would not be a
+        #       good end user experience.
+        #
+        #       Bokeh 3.0 is implementing an "on load" feature, which would nullify this
+        #       requirement, and until that version is released, we have to employ this
+        #       work-around.
+
+        # Create empty Bokeh sources using Python.
+        sources = utils.create_sources(num_chains=self.num_chains)
+
+        # Create empty figures for the tool using Python.
+        figures = utils.create_figures(rv_name=rv_name, num_chains=self.num_chains)
+
+        # Create empty glyphs and attach them to the figures using Python.
+        glyphs = utils.create_glyphs(num_chains=self.num_chains)
+        utils.add_glyphs(sources=sources, figures=figures, glyphs=glyphs)
+
+        # Create empty annotations and attach them to the figures using Python.
+        annotations = utils.create_annotations(
+            figures=figures,
+            num_chains=self.num_chains,
+        )
+        utils.add_annotations(figures=figures, annotations=annotations)
+
+        # Create empty tool tips and attach them to the figures using Python.
+        tooltips = utils.create_tooltips(
+            figures=figures,
+            rv_name=rv_name,
+            num_chains=self.num_chains,
+        )
+        utils.add_tooltips(figures=figures, tooltips=tooltips)
+
+        # Create the widgets for the tool using Python.
+        widgets = utils.create_widgets(rv_names=self.rv_names, rv_name=rv_name)
+
+        # Create callbacks for the tool using JavaScript.
+        callback_js = f"""
+            const rvName = widgets.rv_select.value;
+            const rvData = data[rvName];
+            try {{
+              trace.update(
+                rvData,
+                rvName,
+                bwFactor,
+                hdiProbability,
+                sources,
+                figures,
+                tooltips,
+              );
+            }} catch (error) {{
+              {self.js}
+              trace.update(
+                rvData,
+                rvName,
+                bwFactor,
+                hdiProbability,
+                sources,
+                figures,
+                tooltips,
+              );
+            }}
+        """
+
+        # Each widget requires the following dictionary for the CustomJS method. Notice
+        # that the callback_js object above uses the names defined as keys in the below
+        # object with values defined by the Python objects.
+        callback_arguments = {
+            "data": self.data,
+            "widgets": widgets,
+            "sources": sources,
+            "figures": figures,
+            "tooltips": tooltips,
+        }
+
+        # Each widget requires slightly different JS.
+        rv_select_js = f"""
+            const bwFactor = 1.0;
+            const hdiProbability = 0.89;
+            widgets.bw_factor_slider.value = bwFactor;
+            widgets.hdi_slider.value = 100 * hdiProbability;
+            {callback_js};
+            figures.marginals.reset.emit();
+        """
+        slider_js = f"""
+            const bwFactor = widgets.bw_factor_slider.value;
+            const hdiProbability = widgets.hdi_slider.value / 100;
+            {callback_js};
+        """
+        slider_callback = CustomJS(args=callback_arguments, code=slider_js)
+        rv_select_callback = CustomJS(args=callback_arguments, code=rv_select_js)
+
+        # Tell Python to use the JavaScript.
+        widgets["rv_select"].js_on_change("value", rv_select_callback)
+        widgets["bw_factor_slider"].js_on_change("value", slider_callback)
+        widgets["hdi_slider"].js_on_change("value", slider_callback)
+
+        # Create the view of the tool and serialize it into HTML using static resources
+        # from Bokeh. Embedding the tool in this manner prevents external CDN calls for
+        # JavaScript resources, and prevents the user from having to know where the
+        # Bokeh server is.
+        tool_view = utils.create_view(figures=figures, widgets=widgets)
+        output = file_html(tool_view, resources=INLINE)
+        return output
diff --git a/src/beanmachine/ppl/diagnostics/tools/trace/typing.py b/src/beanmachine/ppl/diagnostics/tools/trace/typing.py
new file mode 100644
index 0000000000..85d9ab5f14
--- /dev/null
+++ b/src/beanmachine/ppl/diagnostics/tools/trace/typing.py
@@ -0,0 +1,158 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Trace diagnostic tool types for a Bean Machine model."""
+
+from typing import Any, Dict, List, Union
+
+from beanmachine.ppl.diagnostics.tools import NotRequired, TypedDict
+from bokeh.models.annotations import Legend
+from bokeh.models.glyphs import Circle, Line, Quad
+from bokeh.models.sources import ColumnDataSource
+from bokeh.models.tools import HoverTool
+from bokeh.models.widgets.inputs import Select
+from bokeh.models.widgets.sliders import Slider
+from bokeh.plotting.figure import Figure
+
+
+# NOTE: These are the types pyre gives us when using `reveal_type(...)` on the outputs
+#       of the methods.
+Data = Dict[str, Dict[Any, Any]]
+Sources = Dict[Any, Any]
+Figures = Dict[Any, Any]
+Glyphs = Dict[Any, Any]
+Annotations = Dict[str, Dict[str, Legend]]
+Tooltips = Dict[Any, Any]
+Widgets = Dict[str, Union[Select, Slider]]
+
+
+# NOTE: TypedDict objects are for reference only. Due to the way pyre accesses keys in
+#       dictionaries, and how NumPy casts arrays when using tolist(), we are unable to
+#       use them, but they provide semantic information for the different types. We must
+#       ignore a lot of lines due to the issue discussed here
+#       https://pyre-check.org/docs/errors/#13-uninitialized-attribute.
+
+
+class _LineOrCircleGlyphData(TypedDict):  # pyre-ignore
+    x: List[float]
+    y: List[float]
+
+
+class _QuadGlyphData(TypedDict):  # pyre-ignore
+    """Follow the RankHistogram interface in stats/histogram.js."""
+
+    left: List[float]
+    top: List[float]
+    right: List[float]
+    bottom: List[float]
+    chain: List[int]
+    draws: List[str]
+    rank: List[float]
+
+
+class _MarginalDataSingleChain(TypedDict):  # pyre-ignore
+    line: _LineOrCircleGlyphData
+    chain: int
+    mean: float
+    bandwidth: float
+
+
+class _ForestDataSingleChain(TypedDict):  # pyre-ignore
+    line: _LineOrCircleGlyphData
+    circle: _LineOrCircleGlyphData
+    chain: int
+    mean: float
+
+
+class _TraceDataSingleChain(TypedDict):  # pyre-ignore
+    line: _LineOrCircleGlyphData
+    chain: int
+    mean: float
+
+
+class _RankDataSingleChain(TypedDict):  # pyre-ignore
+    quad: _QuadGlyphData
+    line: _LineOrCircleGlyphData
+    chain: List[int]
+    rankMean: List[float]
+    mean: List[float]
+
+
+_MarginalDataAllChains = Dict[str, _MarginalDataSingleChain]
+_ForestDataAllChains = Dict[str, _ForestDataSingleChain]
+_TraceDataAllChains = Dict[str, _TraceDataSingleChain]
+_RankDataAllChains = Dict[str, _RankDataSingleChain]
+
+
+class _Data(TypedDict):  # pyre-ignore
+    marginals: _MarginalDataAllChains
+    forests: _ForestDataAllChains
+    traces: _TraceDataAllChains
+    ranks: _RankDataAllChains
+
+
+class _SourceSingleChain(TypedDict):  # pyre-ignore
+    line: ColumnDataSource
+    circle: NotRequired[ColumnDataSource]
+    quad: NotRequired[ColumnDataSource]
+
+
+_SourceAllChains = Dict[str, _SourceSingleChain]
+
+
+class _Sources(TypedDict):  # pyre-ignore
+    marginals: _SourceAllChains
+    forests: _SourceAllChains
+    traces: _SourceAllChains
+    ranks: _SourceAllChains
+
+
+class _Figures(TypedDict):  # pyre-ignore
+    marginals: Figure
+    forests: Figure
+    traces: Figure
+    ranks: Figure
+
+
+class _RankTooltip(TypedDict):  # pyre-ignore
+    line: HoverTool
+    quad: HoverTool
+
+
+class _Tooltips(TypedDict):  # pyre-ignore
+    marginals: List[HoverTool]
+    forests: List[HoverTool]
+    traces: List[HoverTool]
+    ranks: List[_RankTooltip]
+
+
+class _Glyph(TypedDict):  # pyre-ignore
+    glyph: Union[Circle, Line, Quad]
+    hover_glyph: Union[Circle, Line, Quad]
+
+
+class _GlyphSingleChain(TypedDict):  # pyre-ignore
+    line: _Glyph
+    circle: NotRequired[_Glyph]
+    quad: NotRequired[_Glyph]
+
+
+_GlyphAllChains = Dict[str, _GlyphSingleChain]
+
+
+class _Glyphs(TypedDict):  # pyre-ignore
+    marginals: _GlyphAllChains
+    forests: _GlyphAllChains
+    traces: _GlyphAllChains
+    ranks: _GlyphAllChains
+
+
+_Annotations = Dict[str, Dict[str, Legend]]
+
+
+class _Widgets(TypedDict):  # pyre-ignore
+    rv_select: Select
+    bw_factor_slider: Slider
+    hdi_slider: Slider
diff --git a/src/beanmachine/ppl/diagnostics/tools/trace/utils.py b/src/beanmachine/ppl/diagnostics/tools/trace/utils.py
new file mode 100644
index 0000000000..58f95c670a
--- /dev/null
+++ b/src/beanmachine/ppl/diagnostics/tools/trace/utils.py
@@ -0,0 +1,745 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Methods used to generate the diagnostic tool."""
+
+from typing import List
+
+from beanmachine.ppl.diagnostics.tools.trace import typing
+from beanmachine.ppl.diagnostics.tools.utils import plotting_utils
+from bokeh.core.property.wrappers import PropertyValueList
+from bokeh.models.annotations import Legend, LegendItem
+from bokeh.models.glyphs import Circle, Line, Quad
+from bokeh.models.layouts import Column, Row
+from bokeh.models.sources import ColumnDataSource
+from bokeh.models.tools import HoverTool
+from bokeh.models.widgets.inputs import Select
+from bokeh.models.widgets.markups import Div
+from bokeh.models.widgets.panels import Panel, Tabs
+from bokeh.models.widgets.sliders import Slider
+from bokeh.plotting.figure import figure
+
+PLOT_WIDTH = 400
+PLOT_HEIGHT = 500
+TRACE_PLOT_WIDTH = 600
+FIGURE_NAMES = ["marginals", "forests", "traces", "ranks"]
+# Define what the empty data object looks like in order to make the browser handle all
+# computations.
+EMPTY_DATA = {}
+
+
+def create_empty_data(num_chains: int) -> typing.Data:
+    """Create an empty data object for the tool.
+
+    We do not know a priori how many chains a model will have, so we use this method to
+    build an empty data object with the given number of chains.
+
+    Parameters
+    ----------
+    num_chains : int
+        The number of chains from the model.
+
+    Returns
+    -------
+    typing.Data
+        An empty data object to be filled by JavaScript.
+    """
+    output = {
+        "marginals": {},
+        "forests": {},
+        "traces": {},
+        "ranks": {},
+    }
+    for chain in range(num_chains):
+        chain_index = chain + 1
+        chain_name = f"chain{chain_index}"
+        marginal = {
+            "line": {"x": [], "y": []},
+            "chain": [],
+            "mean": [],
+            "bandwidth": [],
+        }
+        forest = {
+            "line": {"x": [], "y": []},
+            "circle": {"x": [], "y": []},
+            "chain": [],
+            "mean": [],
+        }
+        trace = {
+            "line": {"x": [], "y": []},
+            "chain": [],
+            "mean": [],
+        }
+        rank = {
+            "quad": {
+                "left": [],
+                "top": [],
+                "right": [],
+                "bottom": [],
+                "chain": [],
+                "draws": [],
+                "rank": [],
+            },
+            "line": {"x": [], "y": []},
+            "chain": [],
+            "rankMean": [],
+            "mean": [],
+        }
+        single_chain_data = [marginal, forest, trace, rank]
+        chain_data = dict(zip(FIGURE_NAMES, single_chain_data))
+        for figure_name in FIGURE_NAMES:
+            output[figure_name][chain_name] = chain_data[figure_name]
+    return output
+
+
+def create_sources(num_chains: int) -> typing.Sources:
+    """Create Bokeh sources from the given data that will be bound to glyphs.
+
+    Parameters
+    ----------
+    num_chains : int
+        The number of chains from the model.
+
+    Returns
+    -------
+    typing.Sources
+        A dictionary of Bokeh ColumnDataSource objects.
+    """
+    global EMPTY_DATA
+    if not EMPTY_DATA:
+        EMPTY_DATA = create_empty_data(num_chains=num_chains)
+
+    output = {}
+    for figure_name, figure_data in EMPTY_DATA.items():
+        output[figure_name] = {}
+        for chain_name, chain_data in figure_data.items():
+            output[figure_name][chain_name] = {}
+            if figure_name == "marginals":
+                output[figure_name][chain_name]["line"] = ColumnDataSource(
+                    {
+                        "x": chain_data["line"]["x"],
+                        "y": chain_data["line"]["y"],
+                        "chain": chain_data["chain"],
+                        "mean": chain_data["mean"],
+                    },
+                )
+            if figure_name == "forests":
+                output[figure_name][chain_name]["line"] = ColumnDataSource(
+                    {
+                        "x": chain_data["line"]["x"],
+                        "y": chain_data["line"]["y"],
+                    },
+                )
+                output[figure_name][chain_name]["circle"] = ColumnDataSource(
+                    {
+                        "x": chain_data["circle"]["x"],
+                        "y": chain_data["circle"]["y"],
+                        "chain": chain_data["chain"],
+                    },
+                )
+            if figure_name == "traces":
+                output[figure_name][chain_name]["line"] = ColumnDataSource(
+                    {
+                        "x": chain_data["line"]["x"],
+                        "y": chain_data["line"]["y"],
+                        "chain": chain_data["chain"],
+                        "mean": chain_data["mean"],
+                    },
+                )
+            if figure_name == "ranks":
+                output[figure_name][chain_name]["line"] = ColumnDataSource(
+                    {
+                        "x": chain_data["line"]["x"],
+                        "y": chain_data["line"]["y"],
+                        "chain": chain_data["chain"],
+                        "rankMean": chain_data["rankMean"],
+                    },
+                )
+                output[figure_name][chain_name]["quad"] = ColumnDataSource(
+                    {
+                        "left": chain_data["quad"]["left"],
+                        "top": chain_data["quad"]["top"],
+                        "right": chain_data["quad"]["right"],
+                        "bottom": chain_data["quad"]["bottom"],
+                        "chain": chain_data["chain"],
+                        "draws": chain_data["quad"]["draws"],
+                        "rank": chain_data["quad"]["rank"],
+                    },
+                )
+    return output
+
+
+def create_figures(rv_name: str, num_chains: int) -> typing.Figures:
+    """Create the Bokeh figures used for the tool.
+
+    Parameters
+    ----------
+    rv_name : str
+        The string representation of the random variable data.
+    num_chains : int
+        The number of chains from the model.
+
+    Returns
+    -------
+    typing.Figures
+        A dictionary of Bokeh Figure objects.
+    """
+    output = {}
+    for figure_name in FIGURE_NAMES:
+        fig = figure(
+            width=PLOT_WIDTH,
+            height=PLOT_HEIGHT,
+            outline_line_color="black",
+            sizing_mode="scale_both",
+        )
+        plotting_utils.style_figure(fig)
+        # NOTE: There are several figures where we do not want the x-axis to change its
+        #       limits. This is why we set the x_range to an object from Bokeh called
+        #       Range1d.
+        if figure_name == "marginals":
+            fig.title = "Marginal"
+            fig.xaxis.axis_label = rv_name
+            # fig.x_range = Range1d()
+            fig.yaxis.visible = False
+        elif figure_name == "forests":
+            fig.title = "Forest"
+            fig.xaxis.axis_label = rv_name
+            fig.yaxis.axis_label = "Chain"
+            fig.yaxis.minor_tick_line_color = None
+            fig.yaxis.ticker.desired_num_ticks = num_chains
+            # fig.x_range = Range1d()
+        elif figure_name == "traces":
+            fig.title = "Trace"
+            fig.xaxis.axis_label = "Draw from single chain"
+            fig.yaxis.axis_label = rv_name
+            fig.width = TRACE_PLOT_WIDTH
+            # fig.x_range = Range1d()
+        elif figure_name == "ranks":
+            fig.title = "Rank"
+            fig.xaxis.axis_label = "Rank from all chains"
+            fig.yaxis.axis_label = "Chain"
+            fig.width = TRACE_PLOT_WIDTH
+            fig.yaxis.minor_tick_line_color = None
+            fig.yaxis.ticker.desired_num_ticks = num_chains
+        output[figure_name] = fig
+    return output
+
+
+def create_glyphs(num_chains: int) -> typing.Glyphs:
+    """Create the glyphs used for the figures of the tool.
+
+    Parameters
+    ----------
+    num_chains : int
+        The number of chains from the model.
+
+    Returns
+    -------
+    typing.Glyphs
+        A dictionary of Bokeh Glyphs objects.
+    """
+    global EMPTY_DATA
+    if not EMPTY_DATA:
+        EMPTY_DATA = create_empty_data(num_chains=num_chains)
+
+    palette = plotting_utils.choose_palette(num_colors=num_chains)
+    output = {}
+    for figure_name, figure_data in EMPTY_DATA.items():
+        output[figure_name] = {}
+        for i, (chain_name, _) in enumerate(figure_data.items()):
+            output[figure_name][chain_name] = {}
+            color = palette[i]
+            if figure_name == "marginals":
+                output[figure_name][chain_name]["line"] = {
+                    "glyph": Line(
+                        x="x",
+                        y="y",
+                        line_color=color,
+                        line_alpha=0.7,
+                        line_width=2.0,
+                        name=f"{figure_name}{chain_name.title()}LineGlyph",
+                    ),
+                    "hover_glyph": Line(
+                        x="x",
+                        y="y",
+                        line_color=color,
+                        line_alpha=1.0,
+                        line_width=2.0,
+                        name=f"{figure_name}{chain_name.title()}LineHoverGlyph",
+                    ),
+                }
+            elif figure_name == "forests":
+                output[figure_name][chain_name] = {
+                    "line": {
+                        "glyph": Line(
+                            x="x",
+                            y="y",
+                            line_color=color,
+                            line_alpha=0.7,
+                            line_width=2.0,
+                            name=f"{figure_name}{chain_name.title()}LineGlyph",
+                        ),
+                        "hover_glyph": Line(
+                            x="x",
+                            y="y",
+                            line_color=color,
+                            line_alpha=1.0,
+                            line_width=2.0,
+                            name=f"{figure_name}{chain_name.title()}LineHoverGlyph",
+                        ),
+                    },
+                    "circle": {
+                        "glyph": Circle(
+                            x="x",
+                            y="y",
+                            size=10,
+                            fill_color=color,
+                            fill_alpha=0.7,
+                            line_color="white",
+                            name=f"{figure_name}{chain_name.title()}CircleGlyph",
+                        ),
+                        "hover_glyph": Circle(
+                            x="x",
+                            y="y",
+                            size=10,
+                            fill_color=color,
+                            fill_alpha=1.0,
+                            line_color="black",
+                            name=f"{figure_name}{chain_name.title()}CircleHoverGlyph",
+                        ),
+                    },
+                }
+            if figure_name == "traces":
+                output[figure_name][chain_name]["line"] = {
+                    "glyph": Line(
+                        x="x",
+                        y="y",
+                        line_color=color,
+                        line_alpha=0.6,
+                        line_width=0.6,
+                        name=f"{figure_name}{chain_name.title()}LineGlyph",
+                    ),
+                    "hover_glyph": Line(
+                        x="x",
+                        y="y",
+                        line_color=color,
+                        line_alpha=0.6,
+                        line_width=1.0,
+                        name=f"{figure_name}{chain_name.title()}LineHoverGlyph",
+                    ),
+                }
+            if figure_name == "ranks":
+                output[figure_name][chain_name] = {
+                    "quad": {
+                        "glyph": Quad(
+                            left="left",
+                            top="top",
+                            right="right",
+                            bottom="bottom",
+                            fill_color=color,
+                            fill_alpha=0.7,
+                            line_color="white",
+                            name=f"{figure_name}{chain_name.title()}QuadGlyph",
+                        ),
+                        "hover_glyph": Quad(
+                            left="left",
+                            top="top",
+                            right="right",
+                            bottom="bottom",
+                            fill_color=color,
+                            fill_alpha=1.0,
+                            line_color="black",
+                            name=f"{figure_name}{chain_name.title()}QuadHoverGlyph",
+                        ),
+                    },
+                    "line": {
+                        "glyph": Line(
+                            x="x",
+                            y="y",
+                            line_color="grey",
+                            line_alpha=0.7,
+                            line_width=3.0,
+                            line_dash="dashed",
+                            name=f"{figure_name}{chain_name.title()}LineGlyph",
+                        ),
+                        "hover_glyph": Line(
+                            x="x",
+                            y="y",
+                            line_color="grey",
+                            line_alpha=1.0,
+                            line_width=3.0,
+                            line_dash="solid",
+                            name=f"{figure_name}{chain_name.title()}LineGlyph",
+                        ),
+                    },
+                }
+    return output
+
+
+def add_glyphs(
+    figures: typing.Figures,
+    glyphs: typing.Glyphs,
+    sources: typing.Sources,
+) -> None:
+    """Bind source data to glyphs and add the glyphs to the given figures.
+
+    Parameters
+    ----------
+    figures : typing.Figures
+        A dictionary of Bokeh Figure objects.
+    glyphs : typing.Glyphs
+        A dictionary of Bokeh Glyphs objects.
+    sources : typing.Sources
+        A dictionary of Bokeh ColumnDataSource objects.
+
+    Returns
+    -------
+    None
+        Adds data bound glyphs to the given figures directly.
+    """
+    # range_min = []
+    # range_max = []
+    for figure_name, figure_sources in sources.items():
+        fig = figures[figure_name]
+        for chain_name, source in figure_sources.items():
+            chain_glyphs = glyphs[figure_name][chain_name]
+            # NOTE: Every figure has a line glyph, so we always add it here.
+            fig.add_glyph(
+                source_or_glyph=source["line"],
+                glyph=chain_glyphs["line"]["glyph"],
+                hover_glyph=chain_glyphs["line"]["hover_glyph"],
+                name=chain_glyphs["line"]["glyph"].name,
+            )
+            # We want to keep the x-axis from moving when changing queries, so we add
+            # the bounds below from the marginal figure. All figures that need to keep
+            # its range stable are linked to the marginal figure's range below.
+            if figure_name == "marginals":
+                pass
+                # data = source["line"].data["x"]
+                # minimum = min(data) if len(data) != 0 else 0
+                # maximum = max(data) if len(data) != 0 else 1
+                # range_min.append(minimum)
+                # range_max.append(maximum)
+            elif figure_name == "forests":
+                fig.add_glyph(
+                    source_or_glyph=source["circle"],
+                    glyph=chain_glyphs["circle"]["glyph"],
+                    hover_glyph=chain_glyphs["circle"]["hover_glyph"],
+                    name=chain_glyphs["circle"]["glyph"].name,
+                )
+            elif figure_name == "ranks":
+                fig.add_glyph(
+                    source_or_glyph=source["quad"],
+                    glyph=chain_glyphs["quad"]["glyph"],
+                    hover_glyph=chain_glyphs["quad"]["hover_glyph"],
+                    name=chain_glyphs["quad"]["glyph"].name,
+                )
+    # Link figure ranges together.
+    # figures["marginals"].x_range = Range1d(
+    #     start=min(range_min) if len(range_min) != 0 else 0,
+    #     end=max(range_max) if len(range_max) != 0 else 1,
+    # )
+    figures["forests"].x_range = figures["marginals"].x_range
+    # figures["traces"].y_range = figures["marginals"].x_range
+
+
+def create_annotations(figures: typing.Figures, num_chains: int) -> typing.Annotations:
+    """Create any annotations for the figures of the tool.
+
+    Parameters
+    ----------
+    figures : typing.Figures
+        A dictionary of Bokeh Figure objects.
+    num_chains : int
+        The number of chains of the model.
+
+    Returns
+    -------
+    typing.Annotations
+        A dictionary of Bokeh Annotation objects.
+    """
+    renderers = []
+    for _, fig in figures.items():
+        renderers.extend(PropertyValueList(fig.renderers))
+    legend_items = []
+    for chain in range(num_chains):
+        chain_index = chain + 1
+        chain_name = f"chain{chain_index}"
+        legend_items.append(
+            LegendItem(
+                renderers=[
+                    renderer
+                    for renderer in renderers
+                    if chain_name in renderer.name.lower()
+                ],
+                label=chain_name,
+            ),
+        )
+    legend = Legend(
+        items=legend_items,
+        orientation="horizontal",
+        border_line_color="black",
+        click_policy="hide",
+    )
+    output = {"traces": {"legend": legend}, "ranks": {"legend": legend}}
+    return output
+
+
+def add_annotations(figures: typing.Figures, annotations: typing.Annotations) -> None:
+    """Add the given annotations to the given figures of the tool.
+    Parameters
+    ----------
+    figures : typing.Figures
+        A dictionary of Bokeh Figure objects.
+    annotations : typing.Annotations
+        A dictionary of Bokeh Annotation objects.
+    Returns
+    -------
+    None
+        Adds annotations directly to the given figures.
+    """
+    for figure_name, figure_annotations in annotations.items():
+        fig = figures[figure_name]
+        for _, annotation in figure_annotations.items():
+            fig.add_layout(annotation, "below")
+
+
+def create_tooltips(
+    rv_name: str,
+    figures: typing.Figures,
+    num_chains: int,
+) -> typing.Tooltips:
+    """Create hover tools for the glyphs used in the figures of the tool.
+
+    Parameters
+    ----------
+    rv_name : str
+        The string representation of the random variable data.
+    figures : typing.Figures
+        A dictionary of Bokeh Figure objects.
+    num_chains : int
+        The number of chains of the model.
+
+    Returns
+    -------
+    typing.Tooltips
+        A dictionary of Bokeh HoverTools objects.
+    """
+    output = {}
+    for figure_name, fig in figures.items():
+        output[figure_name] = []
+        for chain in range(num_chains):
+            chain_index = chain + 1
+            chain_name = f"chain{chain_index}"
+            if figure_name == "marginals":
+                glyph_name = f"{figure_name}{chain_name.title()}LineGlyph"
+                output[figure_name].append(
+                    HoverTool(
+                        renderers=plotting_utils.filter_renderers(fig, glyph_name),
+                        tooltips=[
+                            ("Chain", "@chain"),
+                            ("Mean", "@mean"),
+                            (rv_name, "@x"),
+                        ],
+                    ),
+                )
+            if figure_name == "forests":
+                glyph_name = f"{figure_name}{chain_name.title()}CircleGlyph"
+                output[figure_name].append(
+                    HoverTool(
+                        renderers=plotting_utils.filter_renderers(fig, glyph_name),
+                        tooltips=[
+                            ("Chain", "@chain"),
+                            (rv_name, "@x"),
+                        ],
+                    ),
+                )
+            if figure_name == "traces":
+                glyph_name = f"{figure_name}{chain_name.title()}LineGlyph"
+                output[figure_name].append(
+                    HoverTool(
+                        renderers=plotting_utils.filter_renderers(fig, glyph_name),
+                        tooltips=[
+                            ("Chain", "@chain"),
+                            ("Mean", "@mean"),
+                            (rv_name, "@y"),
+                        ],
+                    ),
+                )
+            if figure_name == "ranks":
+                output[figure_name].append(
+                    {
+                        "line": HoverTool(
+                            renderers=plotting_utils.filter_renderers(
+                                fig, f"{figure_name}{chain_name.title()}LineGlyph"
+                            ),
+                            tooltips=[("Chain", "@chain"), ("Rank mean", "@rankMean")],
+                        ),
+                        "quad": HoverTool(
+                            renderers=plotting_utils.filter_renderers(
+                                fig, f"{figure_name}{chain_name.title()}QuadGlyph"
+                            ),
+                            tooltips=[
+                                ("Chain", "@chain"),
+                                ("Draws", "@draws"),
+                                ("Rank", "@rank"),
+                            ],
+                        ),
+                    }
+                )
+    return output
+
+
+def add_tooltips(figures: typing.Figures, tooltips: typing.Tooltips) -> None:
+    """Add the given tools to the figures.
+    Parameters
+    ----------
+    figures : typing.Figures
+        A dictionary of Bokeh Figure objects.
+    tooltips : typing.Tooltips
+        A dictionary of Bokeh HoverTools objects.
+    Returns
+    -------
+    None
+        Adds the tooltips directly to the given figures.
+    """
+    for figure_name, fig in figures.items():
+        for tips in tooltips[figure_name]:
+            if figure_name == "ranks":
+                for _, tips_ in tips.items():
+                    fig.add_tools(tips_)
+            else:
+                fig.add_tools(tips)
+
+
+def create_widgets(rv_names: List[str], rv_name: str) -> typing.Widgets:
+    """Create the widgets used in the tool.
+
+    Parameters
+    ----------
+    rv_names : List[str]
+        A list of all available random variable names.
+    rv_name : str
+        The string representation of the random variable data.
+
+    Returns
+    -------
+    typing.Widgets
+        A dictionary of Bokeh widget objects.
+    """
+    output = {
+        "rv_select": Select(value=rv_name, options=rv_names, title="Query"),
+        "bw_factor_slider": Slider(
+            start=0.01,
+            end=2.00,
+            step=0.01,
+            value=1.0,
+            title="Bandwidth factor",
+        ),
+        "hdi_slider": Slider(start=1, end=99, step=1, value=89, title="HDI"),
+    }
+    return output
+
+
+def help_page() -> Div:
+    """Help tab for the tool.
+    Returns
+    -------
+    Div
+        Bokeh Div widget containing the help tab information.
+    """
+    text = """
+    <h2>Rank plots</h2>
+    <p style="margin-bottom: 10px">
+      Rank plots are a histogram of the samples over time. All samples across
+      all chains are ranked and then we plot the average rank for each chain on
+      regular intervals. If the chains are mixing well this histogram should
+      look roughly uniform. If it looks highly irregular that suggests chains
+      might be getting stuck and not adequately exploring the sample space.
+      See the paper by Vehtari <em>et al</em> for more information.
+    </p>
+    <h2>Trace plots</h2>
+    <p style="margin-bottom: 10px">
+      The more familiar trace plots are also included in this widget. You can
+      click on the legend to show/hide different chains and compare them to the
+      rank plots.
+    </p>
+    <ul>
+      <li>
+        Vehtari A, Gelman A, Simpson D, Carpenter B, Bürkner PC (2021)
+        <b>
+          Rank-normalization, folding, and localization: An improved \\(\\hat{R}\\)
+          for assessing convergence of MCMC (with discussion)
+        </b>.
+        <em>Bayesian Analysis</em> 16(2)
+        667–718.
+        <a
+          href=https://dx.doi.org/10.1214/20-BA1221
+          style="color: blue"
+          target="_blank"
+        >
+          doi: 10.1214/20-BA1221
+        </a>.
+      </li>
+    </ul>
+    """
+    return Div(text=text, disable_math=False, min_width=PLOT_WIDTH)
+
+
+def create_view(figures: typing.Figures, widgets: typing.Widgets) -> Tabs:
+    """Create the tool view.
+
+    Parameters
+    ----------
+    figures : typing.Figures
+        A dictionary of Bokeh Figure objects.
+    widgets : typing.Widgets
+        A dictionary of Bokeh widget objects.
+
+    Returns
+    -------
+    Tabs
+        Bokeh Tabs objects.
+    """
+    toolbar = plotting_utils.create_toolbar(list(figures.values()))
+    help_panel = Panel(child=help_page(), title="Help", name="helpPanel")
+    marginal_panel = Panel(
+        child=Column(
+            children=[figures["marginals"], widgets["bw_factor_slider"]],
+            sizing_mode="scale_both",
+        ),
+        title="Marginals",
+    )
+    forest_panel = Panel(
+        child=Column(
+            children=[figures["forests"], widgets["hdi_slider"]],
+            sizing_mode="scale_both",
+        ),
+        title="HDIs",
+    )
+    left_panels = Tabs(tabs=[marginal_panel, forest_panel], sizing_mode="scale_both")
+    trace_panel = Panel(
+        child=Column(children=[figures["traces"]], sizing_mode="scale_both"),
+        title="Traces",
+    )
+    rank_panel = Panel(
+        child=Column(children=[figures["ranks"]], sizing_mode="scale_both"),
+        title="Ranks",
+    )
+    right_panels = Tabs(tabs=[trace_panel, rank_panel], sizing_mode="scale_both")
+    tool_panel = Panel(
+        child=Column(
+            children=[
+                widgets["rv_select"],
+                Row(
+                    children=[left_panels, right_panels, toolbar],
+                    sizing_mode="scale_both",
+                ),
+            ],
+            sizing_mode="scale_both",
+        ),
+        title="Trace tool",
+    )
+    return Tabs(tabs=[tool_panel, help_panel], sizing_mode="scale_both")
diff --git a/src/beanmachine/ppl/diagnostics/tools/utils/base.py b/src/beanmachine/ppl/diagnostics/tools/utils/base.py
index 3a624c219b..6a4e605d8a 100644
--- a/src/beanmachine/ppl/diagnostics/tools/utils/base.py
+++ b/src/beanmachine/ppl/diagnostics/tools/utils/base.py
@@ -7,7 +7,7 @@
 
 import re
 from abc import ABC, abstractmethod
-from typing import TypeVar
+from typing import Optional, TypeVar
 
 from beanmachine.ppl.diagnostics.tools import JS_DIST_DIR
 from beanmachine.ppl.diagnostics.tools.utils import plotting_utils
@@ -38,8 +38,8 @@ def load_js(self: T) -> str:
             js = f.read()
         return js
 
-    def show(self: T):
-        display(HTML(self.create_document()))
+    def show(self: T, name: Optional[str] = None):
+        display(HTML(self.create_document(name=name)))
 
     @abstractmethod
     def create_document(self: T) -> str:
diff --git a/src/beanmachine/ppl/diagnostics/tools/viz.py b/src/beanmachine/ppl/diagnostics/tools/viz.py
index eb3e627c55..d92fcb8f4d 100644
--- a/src/beanmachine/ppl/diagnostics/tools/viz.py
+++ b/src/beanmachine/ppl/diagnostics/tools/viz.py
@@ -53,3 +53,17 @@ def marginal1d(self: DiagnosticsTools) -> None:
         from beanmachine.ppl.diagnostics.tools.marginal1d.tool import Marginal1d
 
         Marginal1d(self.mcs).show()
+
+    def trace(self: T, name: Optional[str] = None) -> None:
+        """
+        Trace diagnostic tool for a Bean Machine model.
+
+        Args:
+            name (:obj:`str`, optional): Optional name for the tool. This is used to
+                persist data as JSON to disk when converting a Jupyter notebook to an
+                MDX file.
+
+        Returns:
+            None: Displays the tool directly in a Jupyter notebook.
+        """
+        Trace(self.mcs).show(name=name)

From dcc83ae2546d97965840cfa2c5286b7dfd0a778a Mon Sep 17 00:00:00 2001
From: Xiaoyan Wang <xxiaoyanwang@gmail.com>
Date: Mon, 17 Oct 2022 21:11:53 -0700
Subject: [PATCH 5/7] Update trace tool to use DiagnosticToolBaseClass

---
 .../ppl/diagnostics/tools/trace/tool.py       | 63 +++++++------------
 .../ppl/diagnostics/tools/utils/base.py       | 46 --------------
 2 files changed, 24 insertions(+), 85 deletions(-)
 delete mode 100644 src/beanmachine/ppl/diagnostics/tools/utils/base.py

diff --git a/src/beanmachine/ppl/diagnostics/tools/trace/tool.py b/src/beanmachine/ppl/diagnostics/tools/trace/tool.py
index 3460c92bd8..7c9555e0b7 100644
--- a/src/beanmachine/ppl/diagnostics/tools/trace/tool.py
+++ b/src/beanmachine/ppl/diagnostics/tools/trace/tool.py
@@ -5,54 +5,44 @@
 
 """Trace diagnostic tool for a Bean Machine model."""
 
-from typing import Optional, TypeVar
+from typing import TypeVar
 
 from beanmachine.ppl.diagnostics.tools.trace import utils
-from beanmachine.ppl.diagnostics.tools.utils.base import Base
+from beanmachine.ppl.diagnostics.tools.utils.diagnostic_tool_base import (
+    DiagnosticToolBaseClass,
+)
 from beanmachine.ppl.inference.monte_carlo_samples import MonteCarloSamples
-from bokeh.embed import file_html
+from bokeh.models import Model
 from bokeh.models.callbacks import CustomJS
-from bokeh.resources import INLINE
 
 
 T = TypeVar("T", bound="Trace")
 
 
-class Trace(Base):
+class Trace(DiagnosticToolBaseClass):
     """Trace tool.
 
-    Parameters
-    ----------
-    mcs : MonteCarloSamples
-        Bean Machine model object.
-
-    Attributes
-    ----------
-    data : Dict[str, List[List[float]]]
-        JSON serialized version of the Bean Machine model.
-    rv_names : List[str]
-        The list of random variables string names for the given model.
-    num_chains : int
-        The number of chains of the model.
-    num_draws : int
-        The number of draws of the model for each chain.
-    palette : List[str]
-        A list of color values used for the glyphs in the figures. The colors are
-        specifically chosen from the Colorblind palette defined in Bokeh.
-    js : str
-        The JavaScript callbacks needed to render the Bokeh tool independently from
-        a Python server.
-    name : str
-        The name to use when saving Bokeh JSON to disk.
+    Args:
+        mcs (MonteCarloSamples): The return object from running a Bean Machine model.
+
+    Attributes:
+        data (Dict[str, List[List[float]]]): JSON serializable representation of the
+            given `mcs` object.
+        rv_names (List[str]): The list of random variables string names for the given
+            model.
+        num_chains (int): The number of chains of the model.
+        num_draws (int): The number of draws of the model for each chain.
+        palette (List[str]): A list of color values used for the glyphs in the figures.
+            The colors are specifically chosen from the Colorblind palette defined in
+            Bokeh.
+        tool_js (str):The JavaScript callbacks needed to render the Bokeh tool
+            independently from a Python server.
     """
 
     def __init__(self: T, mcs: MonteCarloSamples) -> None:
         super(Trace, self).__init__(mcs)
 
-    def create_document(self: T, name: Optional[str] = None) -> str:
-        if name is not None:
-            self.name = name
-
+    def create_document(self: T) -> Model:
         # Initialize widget values using Python.
         rv_name = self.rv_names[0]
 
@@ -113,7 +103,7 @@ def create_document(self: T, name: Optional[str] = None) -> str:
                 tooltips,
               );
             }} catch (error) {{
-              {self.js}
+              {self.tool_js}
               trace.update(
                 rvData,
                 rvName,
@@ -159,10 +149,5 @@ def create_document(self: T, name: Optional[str] = None) -> str:
         widgets["bw_factor_slider"].js_on_change("value", slider_callback)
         widgets["hdi_slider"].js_on_change("value", slider_callback)
 
-        # Create the view of the tool and serialize it into HTML using static resources
-        # from Bokeh. Embedding the tool in this manner prevents external CDN calls for
-        # JavaScript resources, and prevents the user from having to know where the
-        # Bokeh server is.
         tool_view = utils.create_view(figures=figures, widgets=widgets)
-        output = file_html(tool_view, resources=INLINE)
-        return output
+        return tool_view
diff --git a/src/beanmachine/ppl/diagnostics/tools/utils/base.py b/src/beanmachine/ppl/diagnostics/tools/utils/base.py
deleted file mode 100644
index 6a4e605d8a..0000000000
--- a/src/beanmachine/ppl/diagnostics/tools/utils/base.py
+++ /dev/null
@@ -1,46 +0,0 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-#
-# This source code is licensed under the MIT license found in the
-# LICENSE file in the root directory of this source tree.
-
-"""Base class for diagnostic tools of a Bean Machine model."""
-
-import re
-from abc import ABC, abstractmethod
-from typing import Optional, TypeVar
-
-from beanmachine.ppl.diagnostics.tools import JS_DIST_DIR
-from beanmachine.ppl.diagnostics.tools.utils import plotting_utils
-from beanmachine.ppl.diagnostics.tools.utils.model_serializers import serialize_bm
-from beanmachine.ppl.inference.monte_carlo_samples import MonteCarloSamples
-from IPython.display import display, HTML
-
-
-T = TypeVar("T", bound="Base")
-
-
-class Base(ABC):
-    @abstractmethod
-    def __init__(self: T, mcs: MonteCarloSamples) -> None:
-        self.data = serialize_bm(mcs)
-        self.rv_names = ["Select a random variable..."] + list(self.data.keys())
-        self.num_chains = mcs.num_chains
-        self.num_draws = mcs.get_num_samples()
-        self.palette = plotting_utils.choose_palette(self.num_chains)
-        self.js = self.load_js()
-
-    def load_js(self: T) -> str:
-        name = self.__class__.__name__
-        name_tokens = re.findall(r"[A-Z][^A-Z]*", name)
-        name = "_".join(name_tokens)
-        path = JS_DIST_DIR.joinpath(f"{name.lower()}.js")
-        with path.open() as f:
-            js = f.read()
-        return js
-
-    def show(self: T, name: Optional[str] = None):
-        display(HTML(self.create_document(name=name)))
-
-    @abstractmethod
-    def create_document(self: T) -> str:
-        ...

From 2168998f0927355bd9ca89b135bc34bebded2a80 Mon Sep 17 00:00:00 2001
From: Xiaoyan Wang <xxiaoyanwang@gmail.com>
Date: Mon, 17 Oct 2022 21:31:52 -0700
Subject: [PATCH 6/7] Revert the hacky change for N-d tensors

---
 .../ppl/diagnostics/tools/utils/model_serializers.py            | 2 --
 1 file changed, 2 deletions(-)

diff --git a/src/beanmachine/ppl/diagnostics/tools/utils/model_serializers.py b/src/beanmachine/ppl/diagnostics/tools/utils/model_serializers.py
index 67c2b418e9..58e96dd3a6 100644
--- a/src/beanmachine/ppl/diagnostics/tools/utils/model_serializers.py
+++ b/src/beanmachine/ppl/diagnostics/tools/utils/model_serializers.py
@@ -13,10 +13,8 @@
 def serialize_bm(samples: MonteCarloSamples) -> Dict[str, List[List[float]]]:
     """
     Convert Bean Machine models to a JSON serializable object.
-
     Args:
         samples (MonteCarloSamples): Output of a model from Bean Machine.
-
     Returns
         Dict[str, List[List[float]]]: The JSON serializable object for use in the
             diagnostics tools.

From 5c6019f8ce74eb2e16523be84e1bb8f65970f913 Mon Sep 17 00:00:00 2001
From: Andy Maloney <amaloney@quansight.com>
Date: Tue, 18 Oct 2022 19:18:22 -0500
Subject: [PATCH 7/7] Rebase etc

- Rebased against `main`, and fixed merged conflicts.
- Updated copyright requirement
- Updated how models are serialized.
- Update style based on other pushes to the branch
- Added `__future__` annotations
- updated a few docstrings that needed an update
- updated yarn lock based on changes with the marginal1d tool and the
  trace tool.
- removed the 3rd party histogram requirement, since it is easy to
  calculate a histogram.
- other fixes associated with merging etc.
---
 .../ppl/diagnostics/tools/__init__.py         |  2 +-
 .../tools/js/src/stats/histogram.ts           |  8 ++-
 .../ppl/diagnostics/tools/js/yarn.lock        | 61 +------------------
 .../ppl/diagnostics/tools/marginal1d/tool.py  | 10 +--
 .../diagnostics/tools/marginal1d/typing.py    |  1 -
 .../ppl/diagnostics/tools/marginal1d/utils.py |  1 -
 .../ppl/diagnostics/tools/trace/tool.py       | 24 +++++---
 .../ppl/diagnostics/tools/trace/typing.py     |  1 -
 .../ppl/diagnostics/tools/trace/utils.py      | 24 ++------
 .../ppl/diagnostics/tools/utils/accessor.py   | 19 +++---
 .../tools/utils/diagnostic_tool_base.py       | 18 +++---
 .../tools/utils/model_serializers.py          | 10 ++-
 .../diagnostics/tools/utils/plotting_utils.py |  1 -
 src/beanmachine/ppl/diagnostics/tools/viz.py  | 19 +++---
 14 files changed, 68 insertions(+), 131 deletions(-)

diff --git a/src/beanmachine/ppl/diagnostics/tools/__init__.py b/src/beanmachine/ppl/diagnostics/tools/__init__.py
index 9cd2d21afd..43cad634da 100644
--- a/src/beanmachine/ppl/diagnostics/tools/__init__.py
+++ b/src/beanmachine/ppl/diagnostics/tools/__init__.py
@@ -6,7 +6,6 @@
 # flake8: noqa
 
 """Visual diagnostic tools for Bean Machine models."""
-
 import sys
 from pathlib import Path
 
@@ -16,6 +15,7 @@
     #       accepted, see https://peps.python.org/pep-0655/. This is to follow the
     #       interface objects in JavaScript that allow keys to not be required using ?.
     from typing import TypedDict
+
     from typing_extensions import NotRequired
 else:
     from typing_extensions import NotRequired, TypedDict
diff --git a/src/beanmachine/ppl/diagnostics/tools/js/src/stats/histogram.ts b/src/beanmachine/ppl/diagnostics/tools/js/src/stats/histogram.ts
index 394a632c24..40a4868571 100644
--- a/src/beanmachine/ppl/diagnostics/tools/js/src/stats/histogram.ts
+++ b/src/beanmachine/ppl/diagnostics/tools/js/src/stats/histogram.ts
@@ -1,4 +1,10 @@
-/* import {calculateHistogram} from 'compute-histogram'; */
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
 import {linearRange, numericalSort, shape} from './array';
 import {rankData, scaleToOne} from './dataTransformation';
 import {mean as computeMean} from './pointStatistic';
diff --git a/src/beanmachine/ppl/diagnostics/tools/js/yarn.lock b/src/beanmachine/ppl/diagnostics/tools/js/yarn.lock
index a6ef03786a..7f512032f9 100644
--- a/src/beanmachine/ppl/diagnostics/tools/js/yarn.lock
+++ b/src/beanmachine/ppl/diagnostics/tools/js/yarn.lock
@@ -106,7 +106,7 @@
     "@jridgewell/sourcemap-codec" "^1.4.10"
     "@jridgewell/trace-mapping" "^0.3.9"
 
-"@jridgewell/resolve-uri@3.1.0", "@jridgewell/resolve-uri@^3.0.3":
+"@jridgewell/resolve-uri@3.1.0":
   version "3.1.0"
   resolved "https://registry.yarnpkg.com/@jridgewell/resolve-uri/-/resolve-uri-3.1.0.tgz#2203b118c157721addfe69d47b70465463066d78"
   integrity sha512-F2msla3tad+Mfht5cJq7LSXcdudKTWCVYUgw6pLFOOHSTtZlj6SWNYAp+AhuqLmWdBO2X5hPrLcu8cVP8fy28w==
@@ -201,7 +201,7 @@
   resolved "https://registry.yarnpkg.com/@types/json5/-/json5-0.0.29.tgz#ee28707ae94e11d2b827bcbe5270bcea7f3e71ee"
   integrity sha512-dRLjCWHYg4oaA77cxO64oO+7JwCwnIzkZPdrrC71jQmQtlhM556pwKo5bUzqvZndkVbeFLIIi+9TC40JNF5hNQ==
 
-"@types/node@*", "@types/node@^18.0.4":
+"@types/node@*":
   version "18.8.5"
   resolved "https://registry.yarnpkg.com/@types/node/-/node-18.8.5.tgz#6a31f820c1077c3f8ce44f9e203e68a176e8f59e"
   integrity sha512-Bq7G3AErwe5A/Zki5fdD3O6+0zDChhg671NfPjtIcbtzDNZTv4NPKMRFr7gtYPG7y+B8uTiNK4Ngd9T0FTar6Q==
@@ -218,20 +218,6 @@
   dependencies:
     "@types/jquery" "*"
 
-"@typescript-eslint/eslint-plugin@^5.30.5":
-  version "5.40.0"
-  resolved "https://registry.yarnpkg.com/@typescript-eslint/eslint-plugin/-/eslint-plugin-5.40.0.tgz#0159bb71410eec563968288a17bd4478cdb685bd"
-  integrity sha512-FIBZgS3DVJgqPwJzvZTuH4HNsZhHMa9SjxTKAZTlMsPw/UzpEjcf9f4dfgDJEHjK+HboUJo123Eshl6niwEm/Q==
-  dependencies:
-    "@typescript-eslint/scope-manager" "5.40.0"
-    "@typescript-eslint/type-utils" "5.40.0"
-    "@typescript-eslint/utils" "5.40.0"
-    debug "^4.3.4"
-    ignore "^5.2.0"
-    regexpp "^3.2.0"
-    semver "^7.3.7"
-    tsutils "^3.21.0"
-
 "@typescript-eslint/parser@^5.30.5":
   version "5.40.0"
   resolved "https://registry.yarnpkg.com/@typescript-eslint/parser/-/parser-5.40.0.tgz#432bddc1fe9154945660f67c1ba6d44de5014840"
@@ -250,16 +236,6 @@
     "@typescript-eslint/types" "5.40.0"
     "@typescript-eslint/visitor-keys" "5.40.0"
 
-"@typescript-eslint/type-utils@5.40.0":
-  version "5.40.0"
-  resolved "https://registry.yarnpkg.com/@typescript-eslint/type-utils/-/type-utils-5.40.0.tgz#4964099d0158355e72d67a370249d7fc03331126"
-  integrity sha512-nfuSdKEZY2TpnPz5covjJqav+g5qeBqwSHKBvz7Vm1SAfy93SwKk/JeSTymruDGItTwNijSsno5LhOHRS1pcfw==
-  dependencies:
-    "@typescript-eslint/typescript-estree" "5.40.0"
-    "@typescript-eslint/utils" "5.40.0"
-    debug "^4.3.4"
-    tsutils "^3.21.0"
-
 "@typescript-eslint/types@5.40.0":
   version "5.40.0"
   resolved "https://registry.yarnpkg.com/@typescript-eslint/types/-/types-5.40.0.tgz#8de07e118a10b8f63c99e174a3860f75608c822e"
@@ -278,19 +254,6 @@
     semver "^7.3.7"
     tsutils "^3.21.0"
 
-"@typescript-eslint/utils@5.40.0":
-  version "5.40.0"
-  resolved "https://registry.yarnpkg.com/@typescript-eslint/utils/-/utils-5.40.0.tgz#647f56a875fd09d33c6abd70913c3dd50759b772"
-  integrity sha512-MO0y3T5BQ5+tkkuYZJBjePewsY+cQnfkYeRqS6tPh28niiIwPnQ1t59CSRcs1ZwJJNOdWw7rv9pF8aP58IMihA==
-  dependencies:
-    "@types/json-schema" "^7.0.9"
-    "@typescript-eslint/scope-manager" "5.40.0"
-    "@typescript-eslint/types" "5.40.0"
-    "@typescript-eslint/typescript-estree" "5.40.0"
-    eslint-scope "^5.1.1"
-    eslint-utils "^3.0.0"
-    semver "^7.3.7"
-
 "@typescript-eslint/visitor-keys@5.40.0":
   version "5.40.0"
   resolved "https://registry.yarnpkg.com/@typescript-eslint/visitor-keys/-/visitor-keys-5.40.0.tgz#dd2d38097f68e0d2e1e06cb9f73c0173aca54b68"
@@ -684,11 +647,6 @@ core-js-pure@^3.25.1:
   resolved "https://registry.yarnpkg.com/core-js-pure/-/core-js-pure-3.25.5.tgz#79716ba54240c6aa9ceba6eee08cf79471ba184d"
   integrity sha512-oml3M22pHM+igfWHDfdLVq2ShWmjM2V4L+dQEBs0DWVIqEm9WHCwGAlZ6BmyBQGy5sFrJmcx+856D9lVKyGWYg==
 
-create-require@^1.1.0:
-  version "1.1.1"
-  resolved "https://registry.yarnpkg.com/create-require/-/create-require-1.1.1.tgz#c1d7e8f1e5f6cfc9ff65f9cd352d37348756c333"
-  integrity sha512-dcKFX3jn0MpIaXjisoRvexIJVEKzaq7z2rZKxf+MSr9TkdmHmsU4m2lcLojrj/FHl8mk5VxMmYA+ftRkP/3oKQ==
-
 cross-spawn@^7.0.2, cross-spawn@^7.0.3:
   version "7.0.3"
   resolved "https://registry.yarnpkg.com/cross-spawn/-/cross-spawn-7.0.3.tgz#f73a85b9d5d41d045551c177e2882d4ac85728a6"
@@ -786,11 +744,6 @@ emoji-regex@^9.2.2:
   resolved "https://registry.yarnpkg.com/emoji-regex/-/emoji-regex-9.2.2.tgz#840c8803b0d8047f4ff0cf963176b32d4ef3ed72"
   integrity sha512-L18DaJsXSUk2+42pv8mLs5jJT2hqFkFE4j21wOmgbUqsZ2hL72NsUU785g9RXgo3s0ZNgVl42TiHp3ZtOv/Vyg==
 
-emoji-regex@^9.2.2:
-  version "9.2.2"
-  resolved "https://registry.yarnpkg.com/emoji-regex/-/emoji-regex-9.2.2.tgz#840c8803b0d8047f4ff0cf963176b32d4ef3ed72"
-  integrity sha512-L18DaJsXSUk2+42pv8mLs5jJT2hqFkFE4j21wOmgbUqsZ2hL72NsUU785g9RXgo3s0ZNgVl42TiHp3ZtOv/Vyg==
-
 enhanced-resolve@^5.0.0, enhanced-resolve@^5.10.0:
   version "5.10.0"
   resolved "https://registry.yarnpkg.com/enhanced-resolve/-/enhanced-resolve-5.10.0.tgz#0dc579c3bb2a1032e357ac45b8f3a6f3ad4fb1e6"
@@ -1630,11 +1583,6 @@ js-sdsl@^4.1.4:
   resolved "https://registry.yarnpkg.com/js-tokens/-/js-tokens-4.0.0.tgz#19203fb59991df98e3a287050d4647cdeaf32499"
   integrity sha512-RdJUflcE3cUzKiMqQgsCu06FPu9UdIJO0beYbPhHN4k6apgJtifcoCtT9bcxOpYBtpD2kCM6Sbzg4CausW/PKQ==
 
-"js-tokens@^3.0.0 || ^4.0.0":
-  version "4.0.0"
-  resolved "https://registry.yarnpkg.com/js-tokens/-/js-tokens-4.0.0.tgz#19203fb59991df98e3a287050d4647cdeaf32499"
-  integrity sha512-RdJUflcE3cUzKiMqQgsCu06FPu9UdIJO0beYbPhHN4k6apgJtifcoCtT9bcxOpYBtpD2kCM6Sbzg4CausW/PKQ==
-
 js-yaml@^4.1.0:
   version "4.1.0"
   resolved "https://registry.yarnpkg.com/js-yaml/-/js-yaml-4.1.0.tgz#c1fb65f8f5017901cdd2c951864ba18458a10602"
@@ -2008,11 +1956,6 @@ prelude-ls@^1.2.1:
   resolved "https://registry.yarnpkg.com/prelude-ls/-/prelude-ls-1.2.1.tgz#debc6489d7a6e6b0e7611888cec880337d316396"
   integrity sha512-vkcDPrRZo1QZLbn5RLGPpg/WmIQ65qoWWhcGKf/b5eplkkarX0m9z8ppCat4mlOqUsWpyNuYgO3VRyrYHSzX5g==
 
-prettier@^2.7.1:
-  version "2.7.1"
-  resolved "https://registry.yarnpkg.com/prettier/-/prettier-2.7.1.tgz#e235806850d057f97bb08368a4f7d899f7760c64"
-  integrity sha512-ujppO+MkdPqoVINuDFDRLClm7D78qbDt0/NR+wp5FqEZOoTNAjPHWj17QRhu7geIHJfcNhRk1XVQmF8Bp3ye+g==
-
 proj4@^2.7.5:
   version "2.8.0"
   resolved "https://registry.yarnpkg.com/proj4/-/proj4-2.8.0.tgz#b2cb8f3ccd56d4dcc7c3e46155cd02caa804b170"
diff --git a/src/beanmachine/ppl/diagnostics/tools/marginal1d/tool.py b/src/beanmachine/ppl/diagnostics/tools/marginal1d/tool.py
index e6ce876e1c..1e122cba87 100644
--- a/src/beanmachine/ppl/diagnostics/tools/marginal1d/tool.py
+++ b/src/beanmachine/ppl/diagnostics/tools/marginal1d/tool.py
@@ -4,8 +4,7 @@
 # LICENSE file in the root directory of this source tree.
 
 """Marginal 1D diagnostic tool for a Bean Machine model."""
-
-from typing import TypeVar
+from __future__ import annotations
 
 from beanmachine.ppl.diagnostics.tools.marginal1d import utils
 from beanmachine.ppl.diagnostics.tools.utils.diagnostic_tool_base import (
@@ -16,9 +15,6 @@
 from bokeh.models.callbacks import CustomJS
 
 
-T = TypeVar("T", bound="Marginal1d")
-
-
 class Marginal1d(DiagnosticToolBaseClass):
     """
     Marginal 1D diagnostic tool.
@@ -40,10 +36,10 @@ class Marginal1d(DiagnosticToolBaseClass):
             independently from a Python server.
     """
 
-    def __init__(self: T, mcs: MonteCarloSamples) -> None:
+    def __init__(self: Marginal1d, mcs: MonteCarloSamples) -> None:
         super(Marginal1d, self).__init__(mcs)
 
-    def create_document(self: T) -> Model:
+    def create_document(self: Marginal1d) -> Model:
         # Initialize widget values using Python.
         rv_name = self.rv_names[0]
         bw_factor = 1.0
diff --git a/src/beanmachine/ppl/diagnostics/tools/marginal1d/typing.py b/src/beanmachine/ppl/diagnostics/tools/marginal1d/typing.py
index 881cd308e4..8796cc70e3 100644
--- a/src/beanmachine/ppl/diagnostics/tools/marginal1d/typing.py
+++ b/src/beanmachine/ppl/diagnostics/tools/marginal1d/typing.py
@@ -4,7 +4,6 @@
 # LICENSE file in the root directory of this source tree.
 
 """Marginal 1D diagnostic tool types for a Bean Machine model."""
-
 from typing import Any, Dict, List, Union
 
 from beanmachine.ppl.diagnostics.tools import TypedDict
diff --git a/src/beanmachine/ppl/diagnostics/tools/marginal1d/utils.py b/src/beanmachine/ppl/diagnostics/tools/marginal1d/utils.py
index 6ef4dfebf7..ff148175c3 100644
--- a/src/beanmachine/ppl/diagnostics/tools/marginal1d/utils.py
+++ b/src/beanmachine/ppl/diagnostics/tools/marginal1d/utils.py
@@ -4,7 +4,6 @@
 # LICENSE file in the root directory of this source tree.
 
 """Methods used to generate the diagnostic tool."""
-
 from typing import List
 
 import numpy as np
diff --git a/src/beanmachine/ppl/diagnostics/tools/trace/tool.py b/src/beanmachine/ppl/diagnostics/tools/trace/tool.py
index 7c9555e0b7..710acc38b5 100644
--- a/src/beanmachine/ppl/diagnostics/tools/trace/tool.py
+++ b/src/beanmachine/ppl/diagnostics/tools/trace/tool.py
@@ -4,8 +4,7 @@
 # LICENSE file in the root directory of this source tree.
 
 """Trace diagnostic tool for a Bean Machine model."""
-
-from typing import TypeVar
+from __future__ import annotations
 
 from beanmachine.ppl.diagnostics.tools.trace import utils
 from beanmachine.ppl.diagnostics.tools.utils.diagnostic_tool_base import (
@@ -16,9 +15,6 @@
 from bokeh.models.callbacks import CustomJS
 
 
-T = TypeVar("T", bound="Trace")
-
-
 class Trace(DiagnosticToolBaseClass):
     """Trace tool.
 
@@ -39,10 +35,10 @@ class Trace(DiagnosticToolBaseClass):
             independently from a Python server.
     """
 
-    def __init__(self: T, mcs: MonteCarloSamples) -> None:
+    def __init__(self: Trace, mcs: MonteCarloSamples) -> None:
         super(Trace, self).__init__(mcs)
 
-    def create_document(self: T) -> Model:
+    def create_document(self: Trace) -> Model:
         # Initialize widget values using Python.
         rv_name = self.rv_names[0]
 
@@ -88,10 +84,22 @@ def create_document(self: T) -> Model:
         # Create the widgets for the tool using Python.
         widgets = utils.create_widgets(rv_names=self.rv_names, rv_name=rv_name)
 
+        # Create the view of the tool and serialize it into HTML using static resources
+        # from Bokeh. Embedding the tool in this manner prevents external CDN calls for
+        # JavaScript resources, and prevents the user from having to know where the
+        # Bokeh server is.
+        tool_view = utils.create_view(figures=figures, widgets=widgets)
+
         # Create callbacks for the tool using JavaScript.
         callback_js = f"""
             const rvName = widgets.rv_select.value;
             const rvData = data[rvName];
+            let bw = 0.0;
+            // Remove the CSS classes that dim the tool output on initial load.
+            const toolTab = toolView.tabs[0];
+            const toolChildren = toolTab.child.children;
+            const dimmedComponent = toolChildren[1];
+            dimmedComponent.css_classes = [];
             try {{
               trace.update(
                 rvData,
@@ -125,6 +133,7 @@ def create_document(self: T) -> Model:
             "sources": sources,
             "figures": figures,
             "tooltips": tooltips,
+            "toolView": tool_view,
         }
 
         # Each widget requires slightly different JS.
@@ -149,5 +158,4 @@ def create_document(self: T) -> Model:
         widgets["bw_factor_slider"].js_on_change("value", slider_callback)
         widgets["hdi_slider"].js_on_change("value", slider_callback)
 
-        tool_view = utils.create_view(figures=figures, widgets=widgets)
         return tool_view
diff --git a/src/beanmachine/ppl/diagnostics/tools/trace/typing.py b/src/beanmachine/ppl/diagnostics/tools/trace/typing.py
index 85d9ab5f14..1841bc7c70 100644
--- a/src/beanmachine/ppl/diagnostics/tools/trace/typing.py
+++ b/src/beanmachine/ppl/diagnostics/tools/trace/typing.py
@@ -4,7 +4,6 @@
 # LICENSE file in the root directory of this source tree.
 
 """Trace diagnostic tool types for a Bean Machine model."""
-
 from typing import Any, Dict, List, Union
 
 from beanmachine.ppl.diagnostics.tools import NotRequired, TypedDict
diff --git a/src/beanmachine/ppl/diagnostics/tools/trace/utils.py b/src/beanmachine/ppl/diagnostics/tools/trace/utils.py
index 58f95c670a..df99a5cb0a 100644
--- a/src/beanmachine/ppl/diagnostics/tools/trace/utils.py
+++ b/src/beanmachine/ppl/diagnostics/tools/trace/utils.py
@@ -4,7 +4,6 @@
 # LICENSE file in the root directory of this source tree.
 
 """Methods used to generate the diagnostic tool."""
-
 from typing import List
 
 from beanmachine.ppl.diagnostics.tools.trace import typing
@@ -201,7 +200,6 @@ def create_figures(rv_name: str, num_chains: int) -> typing.Figures:
         if figure_name == "marginals":
             fig.title = "Marginal"
             fig.xaxis.axis_label = rv_name
-            # fig.x_range = Range1d()
             fig.yaxis.visible = False
         elif figure_name == "forests":
             fig.title = "Forest"
@@ -209,13 +207,11 @@ def create_figures(rv_name: str, num_chains: int) -> typing.Figures:
             fig.yaxis.axis_label = "Chain"
             fig.yaxis.minor_tick_line_color = None
             fig.yaxis.ticker.desired_num_ticks = num_chains
-            # fig.x_range = Range1d()
         elif figure_name == "traces":
             fig.title = "Trace"
             fig.xaxis.axis_label = "Draw from single chain"
             fig.yaxis.axis_label = rv_name
             fig.width = TRACE_PLOT_WIDTH
-            # fig.x_range = Range1d()
         elif figure_name == "ranks":
             fig.title = "Rank"
             fig.xaxis.axis_label = "Rank from all chains"
@@ -399,8 +395,6 @@ def add_glyphs(
     None
         Adds data bound glyphs to the given figures directly.
     """
-    # range_min = []
-    # range_max = []
     for figure_name, figure_sources in sources.items():
         fig = figures[figure_name]
         for chain_name, source in figure_sources.items():
@@ -417,11 +411,6 @@ def add_glyphs(
             # its range stable are linked to the marginal figure's range below.
             if figure_name == "marginals":
                 pass
-                # data = source["line"].data["x"]
-                # minimum = min(data) if len(data) != 0 else 0
-                # maximum = max(data) if len(data) != 0 else 1
-                # range_min.append(minimum)
-                # range_max.append(maximum)
             elif figure_name == "forests":
                 fig.add_glyph(
                     source_or_glyph=source["circle"],
@@ -437,12 +426,7 @@ def add_glyphs(
                     name=chain_glyphs["quad"]["glyph"].name,
                 )
     # Link figure ranges together.
-    # figures["marginals"].x_range = Range1d(
-    #     start=min(range_min) if len(range_min) != 0 else 0,
-    #     end=max(range_max) if len(range_max) != 0 else 1,
-    # )
     figures["forests"].x_range = figures["marginals"].x_range
-    # figures["traces"].y_range = figures["marginals"].x_range
 
 
 def create_annotations(figures: typing.Figures, num_chains: int) -> typing.Annotations:
@@ -573,13 +557,15 @@ def create_tooltips(
                     {
                         "line": HoverTool(
                             renderers=plotting_utils.filter_renderers(
-                                fig, f"{figure_name}{chain_name.title()}LineGlyph"
+                                fig,
+                                f"{figure_name}{chain_name.title()}LineGlyph",
                             ),
                             tooltips=[("Chain", "@chain"), ("Rank mean", "@rankMean")],
                         ),
                         "quad": HoverTool(
                             renderers=plotting_utils.filter_renderers(
-                                fig, f"{figure_name}{chain_name.title()}QuadGlyph"
+                                fig,
+                                f"{figure_name}{chain_name.title()}QuadGlyph",
                             ),
                             tooltips=[
                                 ("Chain", "@chain"),
@@ -587,7 +573,7 @@ def create_tooltips(
                                 ("Rank", "@rank"),
                             ],
                         ),
-                    }
+                    },
                 )
     return output
 
diff --git a/src/beanmachine/ppl/diagnostics/tools/utils/accessor.py b/src/beanmachine/ppl/diagnostics/tools/utils/accessor.py
index e4e23f6bc6..9146b3c8aa 100644
--- a/src/beanmachine/ppl/diagnostics/tools/utils/accessor.py
+++ b/src/beanmachine/ppl/diagnostics/tools/utils/accessor.py
@@ -10,17 +10,15 @@
 - `pandas`: https://github.com/pandas-dev/pandas/blob/main/pandas/core/accessor.py
 - `xarray`: https://github.com/pydata/xarray/blob/main/xarray/core/extensions.py
 """
+from __future__ import annotations
 
 import contextlib
 import warnings
-from typing import Callable, TypeVar
+from typing import Callable
 
 from beanmachine.ppl.inference.monte_carlo_samples import MonteCarloSamples
 
 
-T = TypeVar("T", bound="CachedAccessor")
-
-
 class CachedAccessor:
     """
     A descriptor for caching accessors.
@@ -38,11 +36,11 @@ class CachedAccessor:
             object.
     """
 
-    def __init__(self: T, name: str, accessor: object) -> None:
+    def __init__(self: CachedAccessor, name: str, accessor: object) -> None:
         self._name = name
         self._accessor = accessor
 
-    def __get__(self: T, obj: object, cls: object) -> object:
+    def __get__(self: CachedAccessor, obj: object, cls: object) -> object:
         """
         Method to retrieve the accessor namespace.
 
@@ -123,7 +121,8 @@ def register_mcs_accessor(name: str) -> Callable:
             object.
 
     Example:
-        >>> from typing import Dict, List, TypeVar
+        >>> from __future__ import annotations
+        >>> from typing import Dict, List
         >>>
         >>> import beanmachine.ppl as bm
         >>> import numpy as np
@@ -132,8 +131,6 @@ def register_mcs_accessor(name: str) -> Callable:
         >>> from beanmachine.ppl.diagnostics.tools.utils import accessor
         >>> from torch import tensor
         >>>
-        >>> T = TypeVar("T", bound="MagicAccessor")
-        >>>
         >>> @bm.random_variable
         >>> def alpha():
         >>>     return dist.Normal(0, 1)
@@ -144,9 +141,9 @@ def register_mcs_accessor(name: str) -> Callable:
         >>>
         >>> @accessor.register_mcs_accessor("magic")
         >>> class MagicAccessor:
-        >>>     def __init__(self: T, mcs: MonteCarloSamples) -> None:
+        >>>     def __init__(self: MagicAccessor, mcs: MonteCarloSamples) -> None:
         >>>         self.mcs = mcs
-        >>>     def show_me(self: T) -> Dict[str, List[List[float]]]:
+        >>>     def show_me(self: MagicAccessor) -> Dict[str, List[List[float]]]:
         >>>         # Return a JSON serializable object from a MonteCarloSamples object.
         >>>         return dict(
         >>>             sorted(
diff --git a/src/beanmachine/ppl/diagnostics/tools/utils/diagnostic_tool_base.py b/src/beanmachine/ppl/diagnostics/tools/utils/diagnostic_tool_base.py
index 9d79084b10..d1c3c0488d 100644
--- a/src/beanmachine/ppl/diagnostics/tools/utils/diagnostic_tool_base.py
+++ b/src/beanmachine/ppl/diagnostics/tools/utils/diagnostic_tool_base.py
@@ -4,10 +4,11 @@
 # LICENSE file in the root directory of this source tree.
 
 """Base class for diagnostic tools of a Bean Machine model."""
+from __future__ import annotations
 
 import re
 from abc import ABC, abstractmethod
-from typing import Any, Mapping, TypeVar
+from typing import Any, Mapping
 
 from beanmachine.ppl.diagnostics.tools import JS_DIST_DIR
 from beanmachine.ppl.diagnostics.tools.utils import plotting_utils
@@ -18,9 +19,6 @@
 from bokeh.resources import INLINE
 
 
-T = TypeVar("T", bound="DiagnosticToolBaseClass")
-
-
 class DiagnosticToolBaseClass(ABC):
     """
     Base class for visual diagnostic tools.
@@ -43,7 +41,7 @@ class DiagnosticToolBaseClass(ABC):
     """
 
     @abstractmethod
-    def __init__(self: T, mcs: MonteCarloSamples) -> None:
+    def __init__(self: DiagnosticToolBaseClass, mcs: MonteCarloSamples) -> None:
         self.data = serialize_bm(mcs)
         self.rv_names = ["Select a random variable..."] + list(self.data.keys())
         self.num_chains = mcs.num_chains
@@ -51,7 +49,7 @@ def __init__(self: T, mcs: MonteCarloSamples) -> None:
         self.palette = plotting_utils.choose_palette(self.num_chains)
         self.tool_js = self.load_tool_js()
 
-    def load_tool_js(self: T) -> str:
+    def load_tool_js(self: DiagnosticToolBaseClass) -> str:
         """
         Load the JavaScript for the diagnostic tool.
 
@@ -75,7 +73,7 @@ def load_tool_js(self: T) -> str:
             tool_js = f.read()
         return tool_js
 
-    def show(self: T) -> None:
+    def show(self: DiagnosticToolBaseClass) -> None:
         """
         Show the diagnostic tool in the notebook.
 
@@ -97,7 +95,7 @@ def show(self: T) -> None:
         html = file_html(doc, resources=INLINE, template=self.html_template())
         display(HTML(html))
 
-    def html_template(self: T) -> str:
+    def html_template(self: DiagnosticToolBaseClass) -> str:
         """
         HTML template object used to inject CSS styles for Bokeh Applications.
 
@@ -145,11 +143,11 @@ def html_template(self: T) -> str:
         """
 
     @abstractmethod
-    def create_document(self: T) -> Model:
+    def create_document(self: DiagnosticToolBaseClass) -> Model:
         """To be implemented by the inheriting class."""
         ...
 
-    def _tool_json(self: T) -> Mapping[Any, Any]:
+    def _tool_json(self: DiagnosticToolBaseClass) -> Mapping[Any, Any]:
         """
         Debugging method used primarily when creating a new diagnostic tool.
 
diff --git a/src/beanmachine/ppl/diagnostics/tools/utils/model_serializers.py b/src/beanmachine/ppl/diagnostics/tools/utils/model_serializers.py
index 58e96dd3a6..27f1d1f8be 100644
--- a/src/beanmachine/ppl/diagnostics/tools/utils/model_serializers.py
+++ b/src/beanmachine/ppl/diagnostics/tools/utils/model_serializers.py
@@ -4,7 +4,6 @@
 # LICENSE file in the root directory of this source tree.
 
 """Collection of serializers for the diagnostics tool use."""
-
 from typing import Dict, List
 
 from beanmachine.ppl.inference.monte_carlo_samples import MonteCarloSamples
@@ -25,6 +24,7 @@ def serialize_bm(samples: MonteCarloSamples) -> Dict[str, List[List[float]]]:
         rv_data = samples[rv_identifier]
         rv_shape = rv_data.shape
         num_rv_chains = rv_shape[0]
+        reshaped_data[f"{str(rv_identifier)}"] = []
         for rv_chain in range(num_rv_chains):
             chain_data = rv_data[rv_chain, :]
             chain_shape = chain_data.shape
@@ -40,8 +40,12 @@ def serialize_bm(samples: MonteCarloSamples) -> Dict[str, List[List[float]]]:
                     reshape_dimensions = chain_shape[1]
                 for i, reshape_dimension in enumerate(range(reshape_dimensions)):
                     data = rv_data[rv_chain, :, reshape_dimension].reshape(-1)
-                    reshaped_data[f"{str(rv_identifier)}[{i}]"] = data.tolist()
+                    if f"{str(rv_identifier)}[{i}]" not in reshaped_data:
+                        reshaped_data[f"{str(rv_identifier)}[{i}]"] = []
+                    reshaped_data[f"{str(rv_identifier)}[{i}]"].append(data.tolist())
             elif len(chain_shape) == 1:
-                reshaped_data[f"{str(rv_identifier)}"] = rv_data[rv_chain, :].tolist()
+                reshaped_data[f"{str(rv_identifier)}"].append(
+                    rv_data[rv_chain, :].tolist(),
+                )
     model = dict(sorted(reshaped_data.items(), key=lambda item: item[0]))
     return model
diff --git a/src/beanmachine/ppl/diagnostics/tools/utils/plotting_utils.py b/src/beanmachine/ppl/diagnostics/tools/utils/plotting_utils.py
index 82b48f2a1f..c116ef778a 100644
--- a/src/beanmachine/ppl/diagnostics/tools/utils/plotting_utils.py
+++ b/src/beanmachine/ppl/diagnostics/tools/utils/plotting_utils.py
@@ -4,7 +4,6 @@
 # LICENSE file in the root directory of this source tree.
 
 """Plotting utilities for the diagnostic tools."""
-
 from typing import List
 
 from bokeh.core.property.nullable import Nullable
diff --git a/src/beanmachine/ppl/diagnostics/tools/viz.py b/src/beanmachine/ppl/diagnostics/tools/viz.py
index d92fcb8f4d..c149ba9d7e 100644
--- a/src/beanmachine/ppl/diagnostics/tools/viz.py
+++ b/src/beanmachine/ppl/diagnostics/tools/viz.py
@@ -49,21 +49,24 @@ def __init__(self: DiagnosticsTools, mcs: MonteCarloSamples) -> None:
 
     @_requires_dev_packages
     def marginal1d(self: DiagnosticsTools) -> None:
-        """Marginal 1D tool."""
+        """
+        Marginal 1D diagnostic tool for a Bean Machine model.
+
+        Returns:
+            None: Displays the tool directly in a Jupyter notebook.
+        """
         from beanmachine.ppl.diagnostics.tools.marginal1d.tool import Marginal1d
 
         Marginal1d(self.mcs).show()
 
-    def trace(self: T, name: Optional[str] = None) -> None:
+    @_requires_dev_packages
+    def trace(self: DiagnosticsTools) -> None:
         """
         Trace diagnostic tool for a Bean Machine model.
 
-        Args:
-            name (:obj:`str`, optional): Optional name for the tool. This is used to
-                persist data as JSON to disk when converting a Jupyter notebook to an
-                MDX file.
-
         Returns:
             None: Displays the tool directly in a Jupyter notebook.
         """
-        Trace(self.mcs).show(name=name)
+        from beanmachine.ppl.diagnostics.tools.trace.tool import Trace
+
+        Trace(self.mcs).show()