Repository: leakec/tfc (language: python)

Issue:
Add type hints
## Feature
I've been lagging on type hints since they are still changing in different Python versions. However, recently JAX has dropped support for numpy 1.19, and the `numpy.typing.NDArray` seems to play well with JAX arrays. Opening this to start looking into type hints. 

Merging from main to get type hints started there.



Discussion / hints:
Keeping this list for types:
* Use `numpy.typing.NDArray` for numpy arrays. This works with JAX and numpy.
* Use `Union[int, float, complex]` or `Union[int, float]` for generic number types. The `numbers.Number` works fine with `mypy`, but linting can sometimes throw errors. If the `Union` versions are used, everything works well with `mypy` and with linting.
Similar to #8, I'm pretty busy at work at the moment, so this issue will be tackled slowly.
Created `type_hints` branch for this issue.


Fix patch:
diff --git a/.github/workflows/run_most_examples.yml b/.github/workflows/run_most_examples.yml
index ad0a23e..6f0695a 100644
--- a/.github/workflows/run_most_examples.yml
+++ b/.github/workflows/run_most_examples.yml
@@ -11,7 +11,7 @@ jobs:
       - uses: actions/checkout@v2
       - uses: actions/setup-python@v2
 
-      - run: "sudo apt-get update && sudo apt-get install -y swig gcc g++"
+      - run: "sudo apt-get update && sudo apt-get install -y swig gcc g++ graphviz"
       - run: python -m pip install wheel setuptools numpy pytest
       - run: python setup.py bdist_wheel
       - run: pip install ./dist/*.whl
diff --git a/MANIFEST.in b/MANIFEST.in
index fc037e0..190ed53 100644
--- a/MANIFEST.in
+++ b/MANIFEST.in
@@ -1,2 +1,3 @@
 include src/tfc/utils/BF/BF.h
 include src/tfc/utils/BF/numpy.i
+include src/tfc/py.typed
diff --git a/docs/requirements.txt b/docs/requirements.txt
index 091216d..e389f5e 100644
--- a/docs/requirements.txt
+++ b/docs/requirements.txt
@@ -8,6 +8,7 @@ ipykernel
 numpy
 jax
 jaxlib
+jaxtyping
 matplotlib
 colorama
 graphviz
diff --git a/examples/Carl_Leake_Dissertation/Chapter_2/Example_2_4.py b/examples/Carl_Leake_Dissertation/Chapter_2/Example_2_4.py
index 026fa4f..9ba1af6 100644
--- a/examples/Carl_Leake_Dissertation/Chapter_2/Example_2_4.py
+++ b/examples/Carl_Leake_Dissertation/Chapter_2/Example_2_4.py
@@ -1,13 +1,14 @@
 from tfc.utils import ComponentConstraintGraph
+from tfc.utils.TFCUtils import ComponentConstraintDict
 
 # Create component constraint dictionary
 N = ['u','v','w']
-E = [{'name':'c1','node0':'u','node1':'v'},
-     {'name':'c1','node0':'v','node1':'w'},
-     {'name':'c1','node0':'w','node1':'u'},
-     {'name':'c2','node0':'u','node1':'v'},
-     {'name':'c3','node0':'u','node1':'v'},
-     {'name':'c4','node0':'v','node1':'w'},]
+E = [ComponentConstraintDict(name='c1',node0='u',node1='v'),
+     ComponentConstraintDict(name='c1',node0='v',node1='w'),
+     ComponentConstraintDict(name='c1',node0='w',node1='u'),
+     ComponentConstraintDict(name='c2',node0='u',node1='v'),
+     ComponentConstraintDict(name='c3',node0='u',node1='v'),
+     ComponentConstraintDict(name='c4',node0='v',node1='w'),]
 
 # Create the component constraint graphs
 p = ComponentConstraintGraph(N,E)
diff --git a/requirements.txt b/requirements.txt
index d3b73f7..ff27e28 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,6 +1,7 @@
 numpy
 jax
 jaxlib
+jaxtyping
 annotated-types
 matplotlib
 plotly
@@ -10,6 +11,9 @@ Pillow
 colorama
 graphviz
 yattag
+tqdm
+pandas
+openpyxl
 
 # Optional
 # pdfCropMargins # Used to crop pdfs in PlotlyMakePlot
diff --git a/setup.py b/setup.py
index 20cd8f0..55cbe72 100644
--- a/setup.py
+++ b/setup.py
@@ -47,11 +47,14 @@ def run(self):
     license="MIT",
     packages=find_packages("src"),
     package_dir={"": "src"},
+    package_data={"": ["src/tfc/py.typed"]},
+    include_package_data=True,
     ext_modules=[BF],
     install_requires=[
         "numpy",
         "jax",
         "jaxlib",
+        "jaxtyping",
         "annotated-types",
         "matplotlib",
         "colorama",
diff --git a/src/tfc/mtfc.py b/src/tfc/mtfc.py
index 1ec314b..bbcd357 100644
--- a/src/tfc/mtfc.py
+++ b/src/tfc/mtfc.py
@@ -1,11 +1,23 @@
-from jax.config import config
+from jax._src.config import config
 
 config.update("jax_enable_x64", True)
 
 from copy import copy
 import numpy as onp
 import jax.numpy as np
-from .utils.types import Literal
+from typing import cast
+import numpy.typing as npt
+from .utils.types import (
+    Literal,
+    uint,
+    IntListOrArray,
+    pint,
+    NumberListOrArray,
+    JaxOrNumpyArray,
+    IntArrayLike,
+    Array,
+    Tuple,
+)
 from jax import core, abstract_arrays
 from jax.interpreters import ad, batching, xla
 from jax.lib import xla_client
@@ -26,25 +38,28 @@ class mtfc:
 
     Parameters
     ----------
-    n : list or array-like
-        Number of points to use per-dimension when discretizing the domain. List or array must be same lenght as number of dimensions.
+    n : IntListOrArray
+        Number of points to use per-dimension when discretizing the domain. List or array must be same length as number of dimensions.
 
-    nC : int or list or array-like
+    nC : IntArrayLike
         Number of functions to remove from the beginning of free function linear expansion. This variable is used to account for basis functions that are linearly dependent on support functions used in the construction of the constrained expressions. The constraints for each dimension can be expressed in 1 of 2 ways. Note that a value of -1 is used to indicate no constraints exist for a particular dimension.
 
         1. As an integer. When expressed as an integer, the first nC basis functions are removed from the free function.
         2. As a list or array. When expressed as a list or array, the basis functions corresponding to the numbers given by the list or array are removed from the free function.
 
-    deg : int
+    deg : uint
         Degree of the basis function expansion.
 
-    basis : {"CP","LeP","FS","ELMTanh","ELMSigmoid","ELMSin","ELMSwish","ELMReLU"}, optional
+    dim : pint
+        Number of dimensions in the domain.
+
+    basis : Literal["CP","LeP","FS","ELMTanh","ELMSigmoid","ELMSin","ELMSwish","ELMReLU"], optional
         This optional keyword argument specifies the basis functions to be used. (Default value = "CP")
 
-    x0 : list or array-like
+    x0 : NumberListOrArray
         Specifies the beginning of the DE domain. (Default value = None)
 
-    xf : list or array-like
+    xf : NumberListOrArray
         Specifies the end of the DE domain. (Default value = None)
 
     backend : Literal["C++", "Python"]
@@ -53,13 +68,15 @@ class mtfc:
 
     def __init__(
         self,
-        n,
-        nC,
-        deg,
-        dim=2,
-        basis="CP",
-        x0=None,
-        xf=None,
+        n: IntListOrArray,
+        nC: IntArrayLike,
+        deg: uint,
+        dim: pint = 2,
+        basis: Literal[
+            "CP", "LeP", "FS", "ELMTanh", "ELMSigmoid", "ELMSin", "ELMSwish", "ELMReLU"
+        ] = "CP",
+        x0: NumberListOrArray = [],
+        xf: NumberListOrArray = [],
         backend: Literal["C++", "Python"] = "C++",
     ):
 
@@ -96,7 +113,7 @@ def __init__(
 
         # Set x0 based on user input
         if x0 is None:
-            self.x0 = np.zeros(dim)
+            self.x0 = onp.zeros(dim)
         else:
             if isinstance(x0, np.ndarray):
                 if not x0.flatten().shape[0] == dim:
@@ -107,21 +124,21 @@ def __init__(
                         + str(dim)
                         + "."
                     )
-                self.x0 = x0
+                self.x0 = cast(onp.ndarray, x0)
             else:
                 if not len(x0) == dim:
                     TFCPrint.Error(
                         "x0 has length "
-                        + len(x0)
+                        + str(len(x0))
                         + ", but it should be equal to the number of dimensions, "
                         + str(dim)
                         + "."
                     )
-                self.x0 = np.array(x0).flatten()
+                self.x0 = onp.array(x0).flatten()
                 if not self.x0.shape[0] == dim:
                     TFCPrint.Error(
                         "x0 has length "
-                        + str(x0.flatten().shape[0])
+                        + str(self.x0.shape[0])
                         + ", but it should be equal to the number of dimensions, "
                         + str(dim)
                         + "."
@@ -129,7 +146,7 @@ def __init__(
 
         # Set xf based on user input
         if xf is None:
-            self.xf = np.zeros(dim)
+            self.xf = onp.zeros(dim)
         else:
             if isinstance(xf, np.ndarray):
                 if not xf.flatten().shape[0] == dim:
@@ -140,21 +157,21 @@ def __init__(
                         + str(dim)
                         + "."
                     )
-                self.xf = xf
+                self.xf = cast(onp.ndarray, xf)
             else:
                 if not len(xf) == dim:
                     TFCPrint.Error(
                         "xf has length "
-                        + len(xf)
+                        + str(len(xf))
                         + ", but it should be equal to the number of dimensions, "
                         + str(dim)
                         + "."
                     )
-                self.xf = np.array(xf).flatten()
+                self.xf = onp.array(xf).flatten()
                 if not self.xf.shape[0] == dim:
                     TFCPrint.Error(
                         "xf has length "
-                        + str(xf.flatten().shape[0])
+                        + str(self.xf.shape[0])
                         + ", but it should be equal to the number of dimensions, "
                         + str(dim)
                         + "."
@@ -165,9 +182,9 @@ def __init__(
             if isinstance(nC, int):
                 self.nC = onp.arange(nC, dtype=onp.int32)
             elif isinstance(nC, np.ndarray):
-                self.nC = nC.astype(onp.int32)
+                self.nC = cast(onp.ndarray, nC.astype(onp.int32))
             elif isinstance(nC, list):
-                self.nC = np.array(nC, dtype=np.int32)
+                self.nC = onp.array(nC, dtype=np.int32)
             if self.nC.shape[0] > self.deg:
                 TFCPrint.Error("Number of basis functions is less than number of constraints!")
             if np.any(self.nC < 0):
@@ -175,6 +192,11 @@ def __init__(
                     "To set nC to -1 (no constraints) either use nC = -1 or nC = 0 (i.e., use an integer not a list or array). Do not put only -1 in a list or array, this will cause issues in the C++ layer."
                 )
         else:
+            if isinstance(nC, int):
+                TFCPrint.Error(
+                    "Cannot use type int for nC when specifying non-ELM type basis function."
+                )
+            # Using explicit type casts here to keep LSPs happy. At this point, we know nC is not a regular integer.
             if isinstance(nC, np.ndarray) and len(nC.shape) > 1:
                 if not nC.shape[0] == self.dim:
                     TFCPrint.Error(
@@ -184,24 +206,24 @@ def __init__(
                         + str(dim)
                         + "."
                     )
-                self.nC = nC.astype(np.int32)
+                self.nC = cast(onp.ndarray, nC.astype(np.int32))
             else:
                 if isinstance(nC, np.ndarray):
                     nC = nC.tolist()
-                if not len(nC) == dim:
+                if not len(cast(IntListOrArray, nC)) == dim:
                     TFCPrint.Error(
                         "nC has length "
-                        + str(len(nC))
+                        + str(len(cast(IntListOrArray, nC)))
                         + ", but it should be equal to the number of dimensions, "
                         + str(dim)
                         + "."
                     )
                 nCmax = 0
                 for k in range(dim):
-                    if isinstance(nC[k], np.ndarray):
-                        nCk = np.array(nC[k]).flatten()
+                    if isinstance(cast(IntListOrArray, nC)[k], np.ndarray):
+                        nCk = np.array(cast(IntListOrArray, nC)[k]).flatten()
                     else:
-                        nCk = np.array([nC[k]]).flatten()
+                        nCk = np.array([cast(IntListOrArray, nC)[k]]).flatten()
                     if nCk.shape[0] == 1:
                         maxk = nCk[0]
                     else:
@@ -216,21 +238,21 @@ def __init__(
 
                 onC = onp.zeros((dim, nCmax))
                 for k in range(dim):
-                    if isinstance(nC[k], np.ndarray):
-                        nCk = np.array(nC[k]).flatten()
+                    if isinstance(cast(IntListOrArray, nC)[k], np.ndarray):
+                        nCk = np.array(cast(IntListOrArray, nC)[k]).flatten()
                     else:
-                        nCk = onp.array([nC[k]]).flatten()
-                    n = nCk.shape[0]
-                    if n == 1:
+                        nCk = onp.array([cast(IntListOrArray, nC)[k]]).flatten()
+                    j = nCk.shape[0]
+                    if j == 1:
                         nCk = onp.arange(nCk[0])
-                        n = nCk.shape[0]
-                    if n < nCmax:
-                        if n == 0:
+                        j = nCk.shape[0]
+                    if j < nCmax:
+                        if j == 0:
                             nCk = -1.0 * onp.ones(nCmax)
                         else:
-                            nCk = np.hstack([nCk, -1 * np.ones(nCmax - n)])
+                            nCk = cast(npt.NDArray, np.hstack([nCk, -1 * np.ones(nCmax - j)]))
                     onC[k, :] = nCk.astype(np.int32)
-                self.nC = np.array(onC.tolist(), dtype=np.int32)
+                self.nC = onp.array(onC.tolist(), dtype=np.int32)
 
         # Setup the basis function
         if backend == "C++":
@@ -285,16 +307,16 @@ def __init__(
 
         # Calculate z points and corresponding x
         self.z = onp.zeros((self.dim, self.N), dtype=self.x0.dtype)
-        x = tuple([onp.zeros(self.N, dtype=self.x0.dtype) for x in range(self.dim)])
+        x = tuple([onp.zeros(self.N, dtype=self.x0.dtype) for _ in range(self.dim)])
         if self.basis in ["CP", "LeP"]:
             for k in range(self.dim):
                 nProd = int(onp.prod(self.n[k + 1 :]))
                 nStack = int(onp.prod(self.n[0:k]))
-                n = self.n[k] - 1
+                j = self.n[k] - 1
                 # Multiplying x0 by 0 here so the array will have the
                 # same type as x0.
-                I = onp.linspace(0 * x0[0], n, n + 1).reshape((n + 1, 1))
-                dark = onp.cos(np.pi * (n - I) / float(n))
+                I = onp.linspace(0 * x0[0], j, j + 1).reshape((j + 1, 1))
+                dark = onp.cos(np.pi * (j - I) / float(j))
                 dark = onp.hstack([dark] * nProd).flatten()
                 self.z[k, :] = onp.array([dark] * nStack).flatten()
                 x[k][:] = (self.z[k, :] - z0) / self.c[k] + self.x0[k]
@@ -302,23 +324,27 @@ def __init__(
             for k in range(self.dim):
                 nProd = int(onp.prod(self.n[k + 1 :]))
                 nStack = int(onp.prod(self.n[0:k]))
-                dark = onp.linspace(z0, zf, num=self.n[k], dtype=x0.dtype).reshape((self.n[k], 1))
+                dark = onp.linspace(z0, zf, num=self.n[k], dtype=self.x0.dtype).reshape(
+                    (self.n[k], 1)
+                )
                 dark = onp.hstack([dark] * nProd).flatten()
                 self.z[k, :] = onp.array([dark] * nStack).flatten()
                 x[k][:] = (self.z[k, :] - z0) / self.c[k] + self.x0[k]
 
-        self.z = np.array(self.z.tolist())
-        self.x = tuple([np.array(x[k].tolist()) for k in range(self.dim)])
+        self.z: Array = cast(Array, np.array(self.z.tolist()))
+        self.x: Tuple[Array, ...] = tuple(
+            [cast(Array, np.array(x[k].tolist())) for k in range(self.dim)]
+        )
 
         self.SetupJAX()
 
-    def H(self, *x, full=False):
+    def H(self, *x: JaxOrNumpyArray, full: bool = False) -> npt.NDArray:
         """
         This function computes the basis function matrix for the points specified by *x.
 
         Parameters
         ----------
-        *x : iterable of array-like
+        *x : JaxOrNumpyArray
             Points to calculate the basis functions at.
 
         full : bool, optional
@@ -326,19 +352,19 @@ def H(self, *x, full=False):
 
         Returns
         -------
-        H : array-like
+        H : NDArray
             Basis function matrix.
         """
         d = onp.zeros(self.dim, dtype=np.int32)
         return self._Hjax(*x, d=d, full=full)
 
-    def Hx(self, *x, full=False):
+    def Hx(self, *x: JaxOrNumpyArray, full: bool = False) -> npt.NDArray:
         """
         This function computes the derivative of the basis function matrix for the points specified by *x with respect to the first variable.
 
         Parameters
         ----------
-        *x : iterable of array-like
+        *x : JaxOrNumpyArray
             Points to calculate the basis functions at.
 
         full : bool, optional
@@ -346,20 +372,20 @@ def Hx(self, *x, full=False):
 
         Returns
         -------
-        Hx : array-like
+        Hx : NDArray
             Derivative of the basis function matrix with respect to the first variable.
         """
         d = onp.zeros(self.dim, dtype=np.int32)
         d[0] = 1
         return self._Hjax(*x, d=d, full=full)
 
-    def Hx2(self, *x, full=False):
+    def Hx2(self, *x: JaxOrNumpyArray, full: bool = False) -> npt.NDArray:
         """
         This function computes the second derivative of the basis function matrix for the points specified by *x with respect to the first variable.
 
         Parameters
         ----------
-        *x : iterable of array-like
+        *x : JaxOrNumpyArray
             Points to calculate the basis functions at.
 
         full : bool, optional
@@ -367,20 +393,20 @@ def Hx2(self, *x, full=False):
 
         Returns
         -------
-        Hx2 : array-like
+        Hx2 : NDArray
             Second derivative of the basis function matrix with respect to the first variable.
         """
         d = onp.zeros(self.dim, dtype=np.int32)
         d[0] = 2
         return self._Hjax(*x, d=d, full=full)
 
-    def Hy2(self, *x, full=False):
+    def Hy2(self, *x: JaxOrNumpyArray, full: bool = False) -> npt.NDArray:
         """
         This function computes the second derivative of the basis function matrix for the points specified by *x with respect to the second variable.
 
         Parameters
         ----------
-        *x : iterable of array-like
+        *x : JaxOrNumpyArray
             Points to calculate the basis functions at.
 
         full : bool, optional
@@ -388,21 +414,21 @@ def Hy2(self, *x, full=False):
 
         Returns
         -------
-        Hy2 : array-like
+        Hy2 : NDArray
             Second derivative of the basis function matrix with respect to the second variable.
         """
         d = onp.zeros(self.dim, dtype=np.int32)
         d[1] = 2
         return self._Hjax(*x, d=d, full=full)
 
-    def Hx2y(self, *x, full=False):
+    def Hx2y(self, *x: JaxOrNumpyArray, full: bool = False) -> npt.NDArray:
         """
         This function computes the mixed derivative (second order derivative with respect to the first variable and first order with respect
         to the second variable) of the basis function matrix for the points specified by *x.
 
         Parameters
         ----------
-        *x : iterable of array-like
+        *x : JaxOrNumpyArray
             Points to calculate the basis functions at.
 
         full : bool, optional
@@ -410,7 +436,7 @@ def Hx2y(self, *x, full=False):
 
         Returns
         -------
-        Hx2y : array-like
+        Hx2y : NDArray
             Mixed derivative of the basis function matrix with respect to the first variable.
         """
         d = onp.zeros(self.dim, dtype=np.int32)
@@ -418,13 +444,13 @@ def Hx2y(self, *x, full=False):
         d[1] = 1
         return self._Hjax(*x, d=d, full=full)
 
-    def Hy(self, *x, full=False):
+    def Hy(self, *x: JaxOrNumpyArray, full: bool = False) -> npt.NDArray:
         """
         This function computes the derivative of the basis function matrix for the points specified by *x with respect to the second variable.
 
         Parameters
         ----------
-        *x : iterable of array-like
+        *x : JaxOrNumpyArray
             Points to calculate the basis functions at.
 
         full : bool, optional
@@ -432,21 +458,21 @@ def Hy(self, *x, full=False):
 
         Returns
         -------
-        Hy : array-like
+        Hy : NDArray
             Derivative of the basis function matrix with respect to the second variable.
         """
         d = onp.zeros(self.dim, dtype=np.int32)
         d[1] = 1
         return self._Hjax(*x, d=d, full=full)
 
-    def Hxy(self, *x, full=False):
+    def Hxy(self, *x: JaxOrNumpyArray, full: bool = False) -> npt.NDArray:
         """
         This function computes the mixed derivative (first order derivative with respect to the first variable and first order with respect
         to the second variable) of the basis function matrix for the points specified by *x.
 
         Parameters
         ----------
-        *x : iterable of array-like
+        *x : JaxOrNumpyArray
             Points to calculate the basis functions at.
 
         full : bool, optional
@@ -454,7 +480,7 @@ def Hxy(self, *x, full=False):
 
         Returns
         -------
-        Hxy : array-like
+        Hxy : NDArray
             Mixed derivative of the basis function matrix with respect to the first variable.
         """
         d = onp.zeros(self.dim, dtype=np.int32)
@@ -462,13 +488,13 @@ def Hxy(self, *x, full=False):
         d[1] = 1
         return self._Hjax(*x, d=d, full=full)
 
-    def Hz(self, *x, full=False):
+    def Hz(self, *x: JaxOrNumpyArray, full: bool = False) -> npt.NDArray:
         """
         This function computes the derivative of the basis function matrix for the points specified by *x with respect to the third variable.
 
         Parameters
         ----------
-        *x : iterable of array-like
+        *x : JaxOrNumpyArray
             Points to calculate the basis functions at.
 
         full : bool, optional
@@ -476,7 +502,7 @@ def Hz(self, *x, full=False):
 
         Returns
         -------
-        Hz : array-like
+        Hz : NDArray
             Derivative of the basis function matrix with respect to the third variable.
         """
         d = onp.zeros(self.dim, dtype=np.int32)
@@ -512,17 +538,19 @@ def _unpack_builder(c):
         # Create Primitives
         H_p = core.Primitive("H")
 
-        def Hjax(*x, d=d0, full=False):
-            return H_p.bind(*x, d=d, full=full)
+        def Hjax(*x: JaxOrNumpyArray, d: npt.NDArray[onp.int32] = d0, full: bool = False):
+            return cast(npt.NDArray, H_p.bind(*x, d=d, full=full))
 
         # Implicit translations
-        def H_impl(*x, d=d0, full=False):
+        def H_impl(*x: npt.NDArray, d: npt.NDArray[onp.int32] = d0, full: bool = False):
             return self.basisClass.H(np.array(x), d, full)
 
         H_p.def_impl(H_impl)
 
         # Define abstract evaluation
-        def H_abstract_eval(*x, d=d0, full=False):
+        def H_abstract_eval(
+            *x, d: npt.NDArray[onp.int32] = d0, full: bool = False
+        ) -> abstract_arrays.ShapedArray:
             if full:
                 dim1 = self.basisClass.numBasisFuncFull
             else:
@@ -537,7 +565,7 @@ def H_abstract_eval(*x, d=d0, full=False):
 
         if self._backend == "C++":
             # XLA compilation
-            def H_xla(c, *x, d=d0, full=False):
+            def H_xla(c, *x, d: npt.NDArray[onp.int32] = d0, full: bool = False):
                 c = _unpack_builder(c)
                 x_shape = c.get_shape(x[0])
                 dims = x_shape.dimensions()
@@ -565,13 +593,13 @@ def H_xla(c, *x, d=d0, full=False):
             xla.backend_specific_translations["cpu"][H_p] = H_xla
 
         # Batching translation
-        def H_batch(vec, batch, d=d0, full=False):
+        def H_batch(vec, batch, d: npt.NDArray[onp.int32] = d0, full: bool = False):
             return Hjax(*vec, d=d, full=full), batch[0]
 
         batching.primitive_batchers[H_p] = H_batch
 
         # Jacobian vector translation
-        def H_jvp(arg_vals, arg_tans, d=d0, full=False):
+        def H_jvp(arg_vals, arg_tans, d: npt.NDArray[onp.int32] = d0, full: bool = False):
             n = len(arg_vals)
             flat = len(arg_vals[0].shape) == 1
             dim0 = arg_vals[0].shape[0]
diff --git a/src/tfc/py.typed b/src/tfc/py.typed
new file mode 100644
index 0000000..e69de29
diff --git a/src/tfc/utfc.py b/src/tfc/utfc.py
index 95cd45c..5e7f9e1 100644
--- a/src/tfc/utfc.py
+++ b/src/tfc/utfc.py
@@ -1,10 +1,12 @@
-from jax.config import config
+from jax._src.config import config
 
 config.update("jax_enable_x64", True)
 
 import numpy as onp
 import jax.numpy as np
-from .utils.types import Literal
+import numpy.typing as npt
+from typing import Optional, cast
+from .utils.types import Literal, uint, IntArrayLike, JaxOrNumpyArray
 from jax import core, abstract_arrays
 from jax.interpreters import ad, batching, xla
 from jax.lib import xla_client
@@ -28,24 +30,33 @@ class utfc:
     ----------
     N : int
         Number of points to use when discretizing the domain.
-    nC : int or list or array-like
+    nC : IntArrayLike
         Number of functions to remove from the free function linear expansion. This variable is used to account for basis functions that are linearly dependent on support functions used in the construction of the constrained expressions. It can be expressed in 1 of 2 ways.
 
         1. As an integer. When expressed as an integer, the first nC basis functions are removed from the free function.
         2. As a list or array. When expressed as a list or array, the basis functions corresponding to the numbers given by the list or array are removed from the free function.
 
-    m : int
+    deg : int
         Degree of the basis function expansion. This number is one less than the number of basis functions used before removing those specified by nC.
     x0 : float, optional
         Specifies the beginning of the DE domain. (Default value = 0)
     xf : float
         This required keyword argument specifies the end of the DE domain.
-    basis : {"CP","LeP","FS","ELMTanh","ELMSigmoid","ELMSin","ELMSwish","ELMReLU"}, optional
+    basis : Literal["CP","LeP","FS","ELMTanh","ELMSigmoid","ELMSin","ELMSwish","ELMReLU"], optional
         This optional keyword argument specifies the basis functions to be used. (Default value = "CP")
     """
 
     def __init__(
-        self, N, nC, deg, basis="CP", x0=None, xf=None, backend: Literal["C++", "Python"] = "C++"
+        self,
+        N: uint,
+        nC: IntArrayLike,
+        deg: uint,
+        basis: Literal[
+            "CP", "LeP", "FS", "ELMTanh", "ELMSigmoid", "ELMSin", "ELMSwish", "ELMReLU"
+        ] = "CP",
+        x0: Optional[float] = None,
+        xf: Optional[float] = None,
+        backend: Literal["C++", "Python"] = "C++",
     ):
         """
         Constructor for the utfc class.
@@ -58,7 +69,7 @@ def __init__(
             Number of functions to remove from the free function linear expansion. This variable is used to account for basis functions that are linearly dependent on support functions used in the construction of the constrained expressions. It can be expressed in 1 of 2 ways.
             1. As an integer. When expressed as an integer, the first nC basis functions are removed from the free function.
             2. As a list or array. When expressed as a list or array, the basis functions corresponding to the numbers given by the list or array are removed from the free function.
-        m : int
+        deg : int
             Degree of the basis function expansion. This number is one less than the number of basis functions used before removing those specified by nC.
         x0 : float, optional
             Specifies the beginning of the DE domain. (Default value = 0)
@@ -76,11 +87,11 @@ def __init__(
         self._backend = backend
 
         if isinstance(nC, int):
-            self.nC = onp.arange(nC, dtype=onp.int32)
+            self.nC: npt.NDArray = onp.arange(nC, dtype=onp.int32)
         elif isinstance(nC, np.ndarray):
-            self.nC = nC.astype(onp.int32)
+            self.nC: npt.NDArray = cast(npt.NDArray, nC.astype(onp.int32))
         elif isinstance(nC, list):
-            self.nC = np.array(nC, dtype=np.int32)
+            self.nC: npt.NDArray = np.array(nC, dtype=np.int32)
         if self.nC.shape[0] > self.deg:
             TFCPrint.Error("Number of basis functions is less than number of constraints!")
         if np.any(self.nC < 0):
@@ -151,7 +162,7 @@ def __init__(
             n = self.N - 1
             # Multiplying x0 by 0 below so the array I has the same
             # type as x0.
-            I = np.linspace(0 * x0, n, n + 1)
+            I = np.linspace(0 * self.x0, n, n + 1)
             self.z = np.cos(np.pi * (n - I) / float(n))
             self.x = (self.z - z0) / self.c + self.x0
         else:
@@ -160,13 +171,13 @@ def __init__(
 
         self._SetupJax()
 
-    def H(self, x, full=False):
+    def H(self, x: JaxOrNumpyArray, full: bool = False) -> npt.NDArray:
         """
         This function computes the basis function matrix for the points specified by x.
 
         Parameters
         ----------
-        x : array-like
+        x : JaxOrNumpyArray
             Points to calculate the basis functions at.
 
         full : bool, optional
@@ -174,17 +185,17 @@ def H(self, x, full=False):
 
         Returns
         -------
-        H : array-like
+        H : NDArray
             Basis function matrix.
         """
         return self._Hjax(x, d=0, full=full)
 
-    def dH(self, x, full=False):
+    def dH(self, x: JaxOrNumpyArray, full: bool = False) -> npt.NDArray:
         """This function computes the deriative of H. See documentation of 'H' for more details.
 
         Parameters
         ----------
-        x : array-like
+        x : JaxOrNumpyArray
             Points to calculate the basis functions at.
 
         full : bool, optional
@@ -192,17 +203,17 @@ def dH(self, x, full=False):
 
         Returns
         -------
-        H : array-like
+        H : NDArray
             Derivative of the basis function matrix.
         """
         return self._Hjax(x, d=1, full=full)
 
-    def d2H(self, x, full=False):
+    def d2H(self, x: JaxOrNumpyArray, full: bool = False) -> npt.NDArray:
         """This function computes the second deriative of H. See documentation of H for more details.
 
         Parameters
         ----------
-        x : array-like
+        x : JaxOrNumpyArray
             Points to calculate the basis functions at.
 
         full : bool, optional
@@ -210,17 +221,17 @@ def d2H(self, x, full=False):
 
         Returns
         -------
-        d2H : array-like
+        d2H : NDArray
             Second derivative of the basis function matrix.
         """
         return self._Hjax(x, d=2, full=full)
 
-    def d4H(self, x, full=False):
+    def d4H(self, x: JaxOrNumpyArray, full: bool = False) -> npt.NDArray:
         """This function computes the fourth deriative of H. See documentation of H for more details.
 
         Parameters
         ----------
-        x : array-like
+        x : JaxOrNumpyArray
             Points to calculate the basis functions at.
 
         full : bool, optional
@@ -228,17 +239,17 @@ def d4H(self, x, full=False):
 
         Returns
         -------
-        d4H : array-like
+        d4H : NDArray
             Fourth derivative of the basis function matrix.
         """
         return self._Hjax(x, d=4, full=full)
 
-    def d8H(self, x, full=False):
+    def d8H(self, x: JaxOrNumpyArray, full: bool = False) -> npt.NDArray:
         """This function computes the eighth deriative of H. See documentation of H for more details.
 
         Parameters
         ----------
-        x : array-like
+        x : JaxOrNumpyArray
             Points to calculate the basis functions at.
 
         full : bool, optional
@@ -246,7 +257,7 @@ def d8H(self, x, full=False):
 
         Returns
         -------
-        d8H : array-like
+        d8H : NDArray
             Eighth derivative of the basis function matrix.
         """
         return self._Hjax(x, d=8, full=full)
@@ -275,17 +286,17 @@ def _unpack_builder(c):
         # Create primitives
         H_p = core.Primitive("H")
 
-        def Hjax(x, d=0, full=False):
-            return H_p.bind(x, d=d, full=full)
+        def Hjax(x: JaxOrNumpyArray, d: uint = 0, full: bool = False) -> npt.NDArray:
+            return cast(npt.NDArray, H_p.bind(x, d=d, full=full))
 
         # Implicit translation
-        def H_impl(x, d=0, full=False):
+        def H_impl(x: npt.NDArray, d: uint = 0, full=False) -> npt.NDArray:
             return self.basisClass.H(x, d, full)
 
         H_p.def_impl(H_impl)
 
         # Abstract evaluation
-        def H_abstract_eval(x, d=0, full=False):
+        def H_abstract_eval(x, d: uint = 0, full: bool = False) -> abstract_arrays.ShapedArray:
             if full:
                 dim1 = self.basisClass.m
             else:
@@ -300,7 +311,7 @@ def H_abstract_eval(x, d=0, full=False):
 
         if self._backend == "C++":
             # XLA compilation
-            def H_xla(c, x, d=0, full=False):
+            def H_xla(c, x, d: uint = 0, full: bool = False):
                 c = _unpack_builder(c)
                 x_shape = c.get_shape(x)
                 dims = x_shape.dimensions()
@@ -327,13 +338,13 @@ def H_xla(c, x, d=0, full=False):
             xla.backend_specific_translations["cpu"][H_p] = H_xla
 
         # Define batching translation
-        def H_batch(vec, batch, d=0, full=False):
+        def H_batch(vec, batch, d: uint = 0, full: bool = False):
             return Hjax(*vec, d=d, full=full), batch[0]
 
         batching.primitive_batchers[H_p] = H_batch
 
         # Define jacobain vector product
-        def H_jvp(arg_vals, arg_tans, d=0, full=False):
+        def H_jvp(arg_vals, arg_tans, d: uint = 0, full: bool = False):
             x = arg_vals[0]
             dx = arg_tans[0]
             if not (dx is ad.Zero):
@@ -386,13 +397,13 @@ def __init__(self, tfcClasses):
             TFCPrint.Error("Not all TFC classes provided have the same number of points.")
         self._tfcClasses = tfcClasses
 
-    def H(self, x, full=False):
+    def H(self, x: JaxOrNumpyArray, full: bool = False) -> npt.NDArray:
         """
         This function computes the basis function matrix for the points specified by x.
 
         Parameters
         ----------
-        x : array-like
+        x : JaxOrNumpyArray
             Points to calculate the basis functions at.
 
         full : bool, optional
@@ -400,18 +411,21 @@ def H(self, x, full=False):
 
         Returns
         -------
-        H : array-like
+        H : NDArray
             Basis function matrix.
         """
-        return np.hstack([k._Hjax(x, d=0, full=full) for j, k in enumerate(self._tfcClasses)])
+        return cast(
+            npt.NDArray,
+            np.hstack([k._Hjax(x, d=0, full=full) for j, k in enumerate(self._tfcClasses)]),
+        )
 
-    def dH(self, x, full=False):
+    def dH(self, x: JaxOrNumpyArray, full: bool = False) -> npt.NDArray:
         """
         This function computes the derivative of the basis function matrix for the points specified by x.
 
         Parameters
         ----------
-        x : array-like
+        x : JaxOrNumpyArray
             Points to calculate the basis functions at.
 
         full : bool, optional
@@ -419,18 +433,21 @@ def dH(self, x, full=False):
 
         Returns
         -------
-        dH : array-like
+        dH : NDArray
             Derivative of the basis function matrix.
         """
-        return np.hstack([k._Hjax(x, d=1, full=full) for j, k in enumerate(self._tfcClasses)])
+        return cast(
+            npt.NDArray,
+            np.hstack([k._Hjax(x, d=1, full=full) for j, k in enumerate(self._tfcClasses)]),
+        )
 
-    def d2H(self, x, full=False):
+    def d2H(self, x: JaxOrNumpyArray, full: bool = False) -> npt.NDArray:
         """
         This function computes the second derivative of the basis function matrix for the points specified by x.
 
         Parameters
         ----------
-        x : array-like
+        x : JaxOrNumpyArray
             Points to calculate the basis functions at.
 
         full : bool, optional
@@ -438,18 +455,21 @@ def d2H(self, x, full=False):
 
         Returns
         -------
-        d2H : array-like
+        d2H : NDArray
             Second derivative of the basis function matrix.
         """
-        return np.hstack([k._Hjax(x, d=2, full=full) for j, k in enumerate(self._tfcClasses)])
+        return cast(
+            npt.NDArray,
+            np.hstack([k._Hjax(x, d=2, full=full) for j, k in enumerate(self._tfcClasses)]),
+        )
 
-    def d3H(self, x, full=False):
+    def d3H(self, x: JaxOrNumpyArray, full: bool = False) -> npt.NDArray:
         """
         This function computes the third derivative of the basis function matrix for the points specified by x.
 
         Parameters
         ----------
-        x : array-like
+        x : JaxOrNumpyArray
             Points to calculate the basis functions at.
 
         full : bool, optional
@@ -457,18 +477,21 @@ def d3H(self, x, full=False):
 
         Returns
         -------
-        d3H : array-like
+        d3H : NDArray
             Third derivative of the basis function matrix.
         """
-        return np.hstack([k._Hjax(x, d=3, full=full) for j, k in enumerate(self._tfcClasses)])
+        return cast(
+            npt.NDArray,
+            np.hstack([k._Hjax(x, d=3, full=full) for j, k in enumerate(self._tfcClasses)]),
+        )
 
-    def d4H(self, x, full=False):
+    def d4H(self, x: JaxOrNumpyArray, full: bool = False) -> npt.NDArray:
         """
         This function computes the fourth derivative of the basis function matrix for the points specified by x.
 
         Parameters
         ----------
-        x : array-like
+        x : JaxOrNumpyArray
             Points to calculate the basis functions at.
 
         full : bool, optional
@@ -476,18 +499,21 @@ def d4H(self, x, full=False):
 
         Returns
         -------
-        d4H : array-like
+        d4H : NDArray
             Fourth derivative of the basis function matrix.
         """
-        return np.hstack([k._Hjax(x, d=4, full=full) for j, k in enumerate(self._tfcClasses)])
+        return cast(
+            npt.NDArray,
+            np.hstack([k._Hjax(x, d=4, full=full) for j, k in enumerate(self._tfcClasses)]),
+        )
 
-    def d8H(self, x, full=False):
+    def d8H(self, x: JaxOrNumpyArray, full: bool = False) -> npt.NDArray:
         """
         This function computes the eighth derivative of the basis function matrix for the points specified by x.
 
         Parameters
         ----------
-        x : array-like
+        x : JaxOrNumpyArray
             Points to calculate the basis functions at.
 
         full : bool, optional
@@ -495,7 +521,10 @@ def d8H(self, x, full=False):
 
         Returns
         -------
-        d8H : array-like
+        d8H : NDArray
             Eighth derivative of the basis function matrix.
         """
-        return np.hstack([k._Hjax(x, d=8, full=full) for j, k in enumerate(self._tfcClasses)])
+        return cast(
+            npt.NDArray,
+            np.hstack([k._Hjax(x, d=8, full=full) for j, k in enumerate(self._tfcClasses)]),
+        )
diff --git a/src/tfc/utils/BF/BF_Py.py b/src/tfc/utils/BF/BF_Py.py
index 4c741bb..ff40fdb 100644
--- a/src/tfc/utils/BF/BF_Py.py
+++ b/src/tfc/utils/BF/BF_Py.py
@@ -3,6 +3,7 @@
 from abc import ABC, abstractmethod
 from numpy import typing as npt
 from tfc.utils.types import uint, Number
+from typing import Callable, Tuple
 
 
 class BasisFunc(ABC):
@@ -28,19 +29,19 @@ def __init__(
         """
         Initialize the basis class.
 
-        Parameters:
-        -----------
-        x0: Number
+        Parameters
+        ----------
+        x0 : Number
             Start of the problem domain.
-        xf: Number
+        xf : Number
             End of the problem domain.
-        nC: npt.NDArray
+        nC : npt.NDArray
             Basis functions to be removed
-        m: uint
+        m : uint
             Number of basis functions.
-        z0: Number
+        z0 : Number
             Start of the basis function domain.
-        zf: Number
+        zf : Number
             End of the basis function domain.
         """
 
@@ -60,19 +61,19 @@ def H(self, x: npt.NDArray, d: uint = 0, full: bool = False) -> npt.NDArray:
         """
         Returns the basis function matrix for the x with a derivative of order d.
 
-        Parameters:
-        -----------
-        x: NDArray
+        Parameters
+        ----------
+        x : NDArray
             Input array. Values to calculate the basis function for.
-        d: uint
+        d : uint
             Order of the derivative
-        full: bool
+        full : bool
             Whether to return the full basis function set, or remove
             the columns associated with self._nC.
 
-        Returns:
-        --------
-        H: NDArray
+        Returns
+        -------
+        H : NDArray
             The basis function values.
         """
 
@@ -90,16 +91,16 @@ def _Hint(self, z: npt.NDArray, d: uint) -> npt.NDArray:
         """
         Internal method used to calcualte the basis function value.
 
-        Parameters:
-        -----------
-        z: NDArray
+        Parameters
+        ----------
+        z : NDArray
             Values to calculate the basis functions for.
-        d: uint
+        d : uint
             Derivative order.
 
-        Returns:
-        --------
-        H: NDArray
+        Returns
+        -------
+        H : NDArray
             Basis function values.
         """
         pass
@@ -110,8 +111,8 @@ def c(self) -> Number:
         Return the constants that map the problem domain to the basis
         function domain.
 
-        Returns:
-        --------
+        Returns
+        -------
         float
             The constant that maps the problem domain to the basis function
             domain.
@@ -135,15 +136,15 @@ def __init__(
         """
         Initialize the basis class.
 
-        Parameters:
-        -----------
-        x0: Number
+        Parameters
+        ----------
+        x0 : Number
             Start of the problem domain.
-        xf: Number
+        xf : Number
             End of the problem domain.
-        nC: npt.NDArray
+        nC : npt.NDArray
             Basis functions to be removed
-        m: uint
+        m:  uint
             Number of basis functions.
         """
         super().__init__(x0, xf, nC, m, -1.0, 1.0)
@@ -152,16 +153,16 @@ def _Hint(self, z: npt.NDArray, d: uint) -> npt.NDArray:
         """
         Internal method used to calcualte the CP basis function values.
 
-        Parameters:
-        -----------
-        z: NDArray
+        Parameters
+        ----------
+        z : NDArray
             Values to calculate the basis functions for.
-        d: uint
+        d : uint
             Derivative order.
 
-        Returns:
-        --------
-        H: NDArray
+        Returns
+        -------
+        H : NDArray
             Basis function values.
         """
         N = np.size(z)
@@ -186,7 +187,10 @@ def _Hint(self, z: npt.NDArray, d: uint) -> npt.NDArray:
             for k in range(2, self._m):
                 F[:, k : k + 1] = 2 * z * F[:, k - 1 : k] - F[:, k - 2 : k - 1]
 
-            def Recurse(dark, d, dCurr=0):
+            def Recurse(dark: npt.NDArray, d: uint, dCurr: uint = 0) -> npt.NDArray:
+                """
+                Take derivative recursively.
+                """
                 if dCurr == d:
                     return dark
                 else:
@@ -221,15 +225,15 @@ def __init__(
         """
         Initialize the basis class.
 
-        Parameters:
-        -----------
-        x0: Number
+        Parameters
+        ----------
+        x0 : Number
             Start of the problem domain.
-        xf: Number
+        xf : Number
             End of the problem domain.
-        nC: npt.NDArray
+        nC : npt.NDArray
             Basis functions to be removed
-        m: uint
+        m : uint
             Number of basis functions.
         """
         super().__init__(x0, xf, nC, m, -1.0, 1.0)
@@ -238,16 +242,16 @@ def _Hint(self, z: npt.NDArray, d: uint) -> npt.NDArray:
         """
         Internal method used to calcualte the LeP basis function values.
 
-        Parameters:
-        -----------
-        z: NDArray
+        Parameters
+        ----------
+        z : NDArray
             Values to calculate the basis functions for.
-        d: uint
+        d : uint
             Derivative order.
 
-        Returns:
-        --------
-        H: NDArray
+        Returns
+        -------
+        H : NDArray
             Basis function values.
         """
         N = np.size(z)
@@ -274,7 +278,10 @@ def _Hint(self, z: npt.NDArray, d: uint) -> npt.NDArray:
                     (2.0 * k + 1.0) * z * F[:, k : k + 1] - k * F[:, k - 1 : k]
                 ) / (k + 1.0)
 
-            def Recurse(dark, d, dCurr=0):
+            def Recurse(dark: npt.NDArray, d: uint, dCurr: uint = 0) -> npt.NDArray:
+                """
+                Take derivative recursively.
+                """
                 if dCurr == d:
                     return dark
                 else:
@@ -303,16 +310,16 @@ def _Hint(self, z: npt.NDArray, d: uint) -> npt.NDArray:
         """
         Internal method used to calcualte the LaP basis function values.
 
-        Parameters:
-        -----------
-        z: NDArray
+        Parameters
+        ----------
+        z : NDArray
             Values to calculate the basis functions for.
-        d: uint
+        d : uint
             Derivative order.
 
-        Returns:
-        --------
-        H: NDArray
+        Returns
+        -------
+        H : NDArray
             Basis function values.
         """
         N = np.size(z)
@@ -339,7 +346,10 @@ def _Hint(self, z: npt.NDArray, d: uint) -> npt.NDArray:
                     (2.0 * k + 1.0 - z) * F[:, k : k + 1] - k * F[:, k - 1 : k]
                 ) / (k + 1.0)
 
-            def Recurse(dark, d, dCurr=0):
+            def Recurse(dark: npt.NDArray, d: uint, dCurr: uint = 0) -> npt.NDArray:
+                """
+                Take derivative recursively.
+                """
                 if dCurr == d:
                     return dark
                 else:
@@ -368,16 +378,16 @@ def _Hint(self, z: npt.NDArray, d: uint) -> npt.NDArray:
         """
         Internal method used to calcualte the HoPpro basis function values.
 
-        Parameters:
-        -----------
-        z: NDArray
+        Parameters
+        ----------
+        z : NDArray
             Values to calculate the basis functions for.
-        d: uint
+        d : uint
             Derivative order.
 
-        Returns:
-        --------
-        H: NDArray
+        Returns
+        -------
+        H : NDArray
             Basis function valuesa
         """
         N = np.size(z)
@@ -402,7 +412,10 @@ def _Hint(self, z: npt.NDArray, d: uint) -> npt.NDArray:
             for k in range(1, self._m - 1):
                 F[:, k + 1 : k + 2] = z * F[:, k : k + 1] - k * F[:, k - 1 : k]
 
-            def Recurse(dark, d, dCurr=0):
+            def Recurse(dark: npt.NDArray, d: uint, dCurr: uint = 0) -> npt.NDArray:
+                """
+                Take derivative recursively.
+                """
                 if dCurr == d:
                     return dark
                 else:
@@ -431,16 +444,16 @@ def _Hint(self, z: npt.NDArray, d: uint) -> npt.NDArray:
         """
         Internal method used to calcualte the HoPpro basis function values.
 
-        Parameters:
-        -----------
-        z: NDArray
+        Parameters
+        ----------
+        z : NDArray
             Values to calculate the basis functions for.
-        d: uint
+        d : uint
             Derivative order.
 
-        Returns:
-        --------
-        H: NDArray
+        Returns
+        -------
+        H : NDArray
             Basis function valuesa
         """
         N = np.size(z)
@@ -465,7 +478,10 @@ def _Hint(self, z: npt.NDArray, d: uint) -> npt.NDArray:
             for k in range(1, self._m - 1):
                 F[:, k + 1 : k + 2] = 2.0 * z * F[:, k : k + 1] - 2.0 * k * F[:, k - 1 : k]
 
-            def Recurse(dark, d, dCurr=0):
+            def Recurse(dark: npt.NDArray, d: uint, dCurr: uint = 0) -> npt.NDArray:
+                """
+                Take derivative recursively.
+                """
                 if dCurr == d:
                     return dark
                 else:
@@ -502,15 +518,15 @@ def __init__(
         """
         Initialize the basis class.
 
-        Parameters:
-        -----------
-        x0: Number
+        Parameters
+        ----------
+        x0 : Number
             Start of the problem domain.
-        xf: Number
+        xf : Number
             End of the problem domain.
-        nC: npt.NDArray
+        nC : npt.NDArray
             Basis functions to be removed
-        m: uint
+        m : uint
             Number of basis functions.
         """
         super().__init__(x0, xf, nC, m, -np.pi, np.pi)
@@ -519,16 +535,16 @@ def _Hint(self, z: npt.NDArray, d: uint) -> npt.NDArray:
         """
         Internal method used to calcualte the CP basis function values.
 
-        Parameters:
-        -----------
-        z: NDArray
+        Parameters
+        ----------
+        z : NDArray
             Values to calculate the basis functions for.
-        d: uint
+        d : uint
             Derivative order.
 
-        Returns:
-        --------
-        H: NDArray
+        Returns
+        -------
+        H : NDArray
             Basis function values.
         """
         N = np.size(z)
@@ -589,15 +605,15 @@ def __init__(
         """
         Initialize the basis class.
 
-        Parameters:
-        -----------
-        x0: Number
+        Parameters
+        ----------
+        x0 : Number
             Start of the problem domain.
-        xf: Number
+        xf : Number
             End of the problem domain.
-        nC: npt.NDArray
+        nC : npt.NDArray
             Basis functions to be removed
-        m: uint
+        m : uint
             Number of basis functions.
         """
         super().__init__(x0, xf, nC, m, 0.0, 1.0)
@@ -614,6 +630,11 @@ def __init__(
     def w(self) -> npt.NDArray:
         """
         Weights of the ELM
+
+        Returns
+        -------
+        NDArray
+            Weights of the ELM.
         """
         return self._w
 
@@ -621,6 +642,11 @@ def w(self) -> npt.NDArray:
     def b(self) -> npt.NDArray:
         """
         Biases of the ELM
+
+        Returns
+        -------
+        NDArray
+            Biases of the ELM.
         """
         return self._b
 
@@ -628,6 +654,11 @@ def b(self) -> npt.NDArray:
     def w(self, val: npt.NDArray) -> None:
         """
         Weights of the ELM.
+
+        Parameters
+        ----------
+        val : NDArray
+            New weights.
         """
         if val.size == self._m:
             self._w = val
@@ -642,6 +673,11 @@ def w(self, val: npt.NDArray) -> None:
     def b(self, val: npt.NDArray) -> None:
         """
         Biases of the ELM.
+
+        Parameters
+        ----------
+        val : NDArray
+            New biases.
         """
         if val.size == self._m:
             self._b = val
@@ -658,16 +694,16 @@ def _Hint(self, z: npt.NDArray, d: uint) -> npt.NDArray:
         """
         Internal method used to calcualte the ELMRelu basis function values.
 
-        Parameters:
-        -----------
-        z: NDArray
+        Parameters
+        ----------
+        z : NDArray
             Values to calculate the basis functions for.
-        d: uint
+        d : uint
             Derivative order.
 
-        Returns:
-        --------
-        H: NDArray
+        Returns
+        -------
+        H : NDArray
             Basis function values.
         """
 
@@ -684,16 +720,16 @@ def _Hint(self, z: npt.NDArray, d: uint) -> npt.NDArray:
         """
         Internal method used to calcualte the ELMSigmoid basis function values.
 
-        Parameters:
-        -----------
-        z: NDArray
+        Parameters
+        ----------
+        z : NDArray
             Values to calculate the basis functions for.
-        d: uint
+        d : uint
             Derivative order.
 
-        Returns:
-        --------
-        H: NDArray
+        Returns
+        -------
+        H : NDArray
             Basis function values.
         """
 
@@ -701,7 +737,12 @@ def _Hint(self, z: npt.NDArray, d: uint) -> npt.NDArray:
 
         f = lambda x: 1.0 / (1.0 + jnp.exp(-self._w * x - self._b))
 
-        def Recurse(dark, d, dCurr=0):
+        def Recurse(
+            dark: Callable[[npt.NDArray], jnp.ndarray], d: uint, dCurr: uint = 0
+        ) -> Callable[[npt.NDArray], jnp.ndarray]:
+            """
+            Take derivative recursively.
+            """
             if dCurr == d:
                 return dark
             else:
@@ -709,7 +750,7 @@ def Recurse(dark, d, dCurr=0):
                 dCurr += 1
                 return Recurse(dark2, d, dCurr=dCurr)
 
-        return Recurse(f, d)(z).to_py()
+        return np.asarray(Recurse(f, d)(z))
 
 
 class ELMTanh(ELM):
@@ -717,16 +758,16 @@ def _Hint(self, z: npt.NDArray, d: uint) -> npt.NDArray:
         """
         Internal method used to calcualte the ELMTanh basis function values.
 
-        Parameters:
-        -----------
-        z: NDArray
+        Parameters
+        ----------
+        z : NDArray
             Values to calculate the basis functions for.
-        d: uint
+        d : uint
             Derivative order.
 
-        Returns:
-        --------
-        H: NDArray
+        Returns
+        -------
+        H : NDArray
             Basis function values.
         """
 
@@ -734,7 +775,12 @@ def _Hint(self, z: npt.NDArray, d: uint) -> npt.NDArray:
 
         f = lambda x: jnp.tanh(self._w * x + self._b)
 
-        def Recurse(dark, d, dCurr=0):
+        def Recurse(
+            dark: Callable[[npt.NDArray], jnp.ndarray], d: uint, dCurr: uint = 0
+        ) -> Callable[[npt.NDArray], jnp.ndarray]:
+            """
+            Take derivative recursively.
+            """
             if dCurr == d:
                 return dark
             else:
@@ -742,7 +788,7 @@ def Recurse(dark, d, dCurr=0):
                 dCurr += 1
                 return Recurse(dark2, d, dCurr=dCurr)
 
-        return Recurse(f, d)(z).to_py()
+        return np.asarray(Recurse(f, d)(z))
 
 
 class ELMSin(ELM):
@@ -750,16 +796,16 @@ def _Hint(self, z: npt.NDArray, d: uint) -> npt.NDArray:
         """
         Internal method used to calcualte the ELMSin basis function values.
 
-        Parameters:
-        -----------
-        z: NDArray
+        Parameters
+        ----------
+        z : NDArray
             Values to calculate the basis functions for.
-        d: uint
+        d : uint
             Derivative order.
 
-        Returns:
-        --------
-        H: NDArray
+        Returns
+        -------
+        H : NDArray
             Basis function values.
         """
 
@@ -767,7 +813,12 @@ def _Hint(self, z: npt.NDArray, d: uint) -> npt.NDArray:
 
         f = lambda x: jnp.sin(self._w * x + self._b)
 
-        def Recurse(dark, d, dCurr=0):
+        def Recurse(
+            dark: Callable[[npt.NDArray], jnp.ndarray], d: uint, dCurr: uint = 0
+        ) -> Callable[[npt.NDArray], jnp.ndarray]:
+            """
+            Take derivative recursively.
+            """
             if dCurr == d:
                 return dark
             else:
@@ -775,7 +826,7 @@ def Recurse(dark, d, dCurr=0):
                 dCurr += 1
                 return Recurse(dark2, d, dCurr=dCurr)
 
-        return Recurse(f, d)(z).to_py()
+        return np.asarray(Recurse(f, d)(z))
 
 
 class ELMSwish(ELM):
@@ -783,16 +834,16 @@ def _Hint(self, z: npt.NDArray, d: uint) -> npt.NDArray:
         """
         Internal method used to calcualte the ELMSwish basis function values.
 
-        Parameters:
-        -----------
-        z: NDArray
+        Parameters
+        ----------
+        z : NDArray
             Values to calculate the basis functions for.
-        d: uint
+        d : uint
             Derivative order.
 
-        Returns:
-        --------
-        H: NDArray
+        Returns
+        -------
+        H : NDArray
             Basis function values.
         """
 
@@ -800,7 +851,12 @@ def _Hint(self, z: npt.NDArray, d: uint) -> npt.NDArray:
 
         f = lambda x: (self._w * x + self._b) / (1.0 + jnp.exp(-self._w * x - self._b))
 
-        def Recurse(dark, d, dCurr=0):
+        def Recurse(
+            dark: Callable[[npt.NDArray], jnp.ndarray], d: uint, dCurr: uint = 0
+        ) -> Callable[[npt.NDArray], jnp.ndarray]:
+            """
+            Take derivative recursively.
+            """
             if dCurr == d:
                 return dark
             else:
@@ -808,7 +864,7 @@ def Recurse(dark, d, dCurr=0):
                 dCurr += 1
                 return Recurse(dark2, d, dCurr=dCurr)
 
-        return Recurse(f, d)(z).to_py()
+        return np.asarray(Recurse(f, d)(z))
 
 
 class nBasisFunc(BasisFunc):
@@ -829,19 +885,19 @@ def __init__(
         """
         Initialize the basis class.
 
-        Parameters:
-        -----------
-        x0: NDArray
+        Parameters
+        ----------
+        x0 : NDArray
             Start of the problem domain.
-        xf: NDArray
+        xf : NDArray
             End of the problem domain.
-        nC: NDArray
+        nC : NDArray
             Basis functions to be removed
-        m: uint
+        m : uint
             Number of basis functions.
-        z0: Number
+        z0 : Number
             Start of the basis function domain.
-        zf: Number
+        zf : Number
             End of the basis function domain.
         """
 
@@ -863,7 +919,26 @@ def __init__(
         self._numBasisFunc = self._NumBasisFunc(self._dim - 1, vec, full=False)
         self._numBasisFuncFull = self._NumBasisFunc(self._dim - 1, vec, full=True)
 
-    def _NumBasisFunc(self, dim: int, vec: npt.NDArray, n: int = 0, full: bool = False):
+    def _NumBasisFunc(self, dim: int, vec: npt.NDArray, n: int = 0, full: bool = False) -> int:
+        """
+        Calculate the number of basis functions.
+
+        Parameters
+        ----------
+        dim : int
+            Number of dimensions.
+        vec : NDArray
+            Vector used to keep track of the order of the basis function.
+        n : int, optional
+            Count of the number of basis functions so far. (Default value = 0)
+        full : bool, optional
+            If true, then does not remove basis functions based on self._nC. (Default value = False)
+
+        Returns
+        -------
+        int
+            Number of basis functions.
+        """
         if dim > 0:
             for x in range(self._m):
                 vec[dim] = x
@@ -890,8 +965,8 @@ def c(self) -> npt.NDArray:
         Return the constants that map the problem domain to the basis
         function domain.
 
-        Returns:
-        --------
+        Returns
+        -------
         npt.NDArray
             The constants that map the problem domain to the basis function
             domain.
@@ -905,8 +980,8 @@ def numBasisFunc(self) -> float:
         Return the number of basis functions once user-specified
         functions have been removed.
 
-        Returns:
-        --------
+        Returns
+        -------
         float:
             The number of basis functions once the user-specified
             functions have been removed.
@@ -920,8 +995,8 @@ def numBasisFuncFull(self) -> float:
         Return the number of basis functions before the user-specified
         functions have been removed.
 
-        Returns:
-        --------
+        Returns
+        -------
         float:
             The number of basis functions before the user-specified
             functions have been removed.
@@ -933,20 +1008,20 @@ def H(self, x: npt.NDArray, d: npt.NDArray, full: bool = False) -> npt.NDArray:
         """
         Returns the basis function matrix for the x with a derivative of order d.
 
-        Parameters:
+        Parameters
         -----------
-        x: NDArray
+        x : NDArray
             Input array. Values to calculate the basis function for.
             Should be size dim x N.
-        d: NDArray
+        d : NDArray
             Order of the derivative
         full: bool
             Whether to return the full basis function set, or remove
             the columns associated with self._nC.
 
-        Returns:
-        --------
-        H: NDArray
+        Returns
+        -------
+        H : NDArray
             The basis function values.
         """
 
@@ -967,12 +1042,50 @@ def H(self, x: npt.NDArray, d: npt.NDArray, full: bool = False) -> npt.NDArray:
 
         # Define functions for use in generating the CP sheet
         def MultT(vec: npt.NDArray) -> npt.NDArray:
+            """
+            Creates basis functions for the multidimensional case by mulitplying the basis functions
+            for the single dimensional cases together.
+
+            Parameters
+            ----------
+            vec : NDArray
+                Used to track the basis functions used from the single dimensional cases.
+
+            Returns
+            -------
+            NDArray
+                Basis functions for the multidimensional case.
+            """
             tout = np.ones((N, 1), dtype=z.dtype)
             for k in range(self._dim):
                 tout *= T[:, vec[k, 0] : vec[k, 0] + 1, k]
             return tout
 
-        def Recurse(dim: int, out: npt.NDArray, vec: npt.NDArray, n: int = 0, full: bool = False):
+        def Recurse(
+            dim: int, out: npt.NDArray, vec: npt.NDArray, n: int = 0, full: bool = False
+        ) -> Tuple[npt.NDArray, int]:
+            """
+            Creates basis functions for the multidimensional case given the basis functions
+            for the single dimensional cases.
+
+            Parameters
+            ----------
+            dim : int
+                Number of dimensions.
+            out : NDArray
+                Basis function for the multidimensional case created so far.
+            n : int, optional
+                Count of the number of basis functions created so far. (Default value = 0)
+            full : bool, optional
+                If true, then does not remove basis functions based on self._nC. (Default value = False)
+
+            Returns
+            -------
+            out : NDArraY
+                Basis functions for the multidimensional case created so far.
+            n : int
+                Basis function count.
+            """
             if dim > 0:
                 for x in range(self._m):
                     vec[dim] = x
@@ -1024,15 +1137,15 @@ def __init__(
         """
         Initialize the n-dimensional CP class.
 
-        Parameters:
-        -----------
-        x0: NDArray
+        Parameters
+        ----------
+        x0 : NDArray
             Start of the problem domain.
-        xf: NDArray
+        xf : NDArray
             End of the problem domain.
-        nC: NDArray
+        nC : NDArray
             Basis functions to be removed
-        m: uint
+        m : uint
             Number of basis functions.
         """
 
@@ -1054,15 +1167,15 @@ def __init__(
         """
         Initialize the n-dimensional LeP class.
 
-        Parameters:
-        -----------
-        x0: NDArray
+        Parameters
+        ----------
+        x0 : NDArray
             Start of the problem domain.
-        xf: NDArray
+        xf : NDArray
             End of the problem domain.
-        nC: NDArray
+        nC : NDArray
             Basis functions to be removed
-        m: uint
+        m : uint
             Number of basis functions.
         """
 
@@ -1084,15 +1197,15 @@ def __init__(
         """
         Initialize the n-dimensional FS class.
 
-        Parameters:
-        -----------
-        x0: NDArray
+        Parameters
+        ----------
+        x0 : NDArray
             Start of the problem domain.
-        xf: NDArray
+        xf : NDArray
             End of the problem domain.
-        nC: NDArray
+        nC : NDArray
             Basis functions to be removed
-        m: uint
+        m : uint
             Number of basis functions.
         """
 
@@ -1116,19 +1229,19 @@ def __init__(
         """
         Initialize the basis class.
 
-        Parameters:
-        -----------
-        x0: NDArray
+        Parameters
+        ----------
+        x0 : NDArray
             Start of the problem domain.
-        xf: NDArray
+        xf : NDArray
             End of the problem domain.
-        nC: NDArray
+        nC : NDArray
             Basis functions to be removed
-        m: uint
+        m : uint
             Number of basis functions.
-        z0: Number
+        z0 : Number
             Start of the basis function domain.
-        zf: Number
+        zf : Number
             End of the basis function domain.
         """
 
@@ -1163,6 +1276,11 @@ def __init__(
     def w(self) -> npt.NDArray:
         """
         Weights of the nELM
+
+        Returns
+        -------
+        NDArray
+            Weights of the ELM.
         """
         return self._w
 
@@ -1170,6 +1288,11 @@ def w(self) -> npt.NDArray:
     def b(self) -> npt.NDArray:
         """
         Biases of the nELM
+
+        Returns
+        -------
+        NDArray
+            Biases of the ELM.
         """
         return self._b
 
@@ -1177,6 +1300,11 @@ def b(self) -> npt.NDArray:
     def w(self, val: npt.NDArray) -> None:
         """
         Weights of the nELM.
+
+        Parameters
+        ----------
+        val : NDArray
+            New weights.
         """
         if val.size == self._m * self._dim:
             self._w = val
@@ -1191,6 +1319,11 @@ def w(self, val: npt.NDArray) -> None:
     def b(self, val: npt.NDArray) -> None:
         """
         Biases of the nELM.
+
+        Parameters
+        ----------
+        val : NDArray
+            New biases.
         """
         if val.size == self._m:
             self._b = val
@@ -1205,20 +1338,20 @@ def H(self, x: npt.NDArray, d: npt.NDArray, full: bool = False) -> npt.NDArray:
         """
         Returns the basis function matrix for the x with a derivative of order d.
 
-        Parameters:
-        -----------
-        x: NDArray
+        Parameters
+        ----------
+        x : NDArray
             Input array. Values to calculate the basis function for.
             Should be size dim x N.
-        d: NDArray
+        d : NDArray
             Order of the derivative
-        full: bool
+        full : bool
             Whether to return the full basis function set, or remove
             the columns associated with self._nC.
 
-        Returns:
-        --------
-        H: NDArray
+        Returns
+        -------
+        H : NDArray
             The basis function values.
         """
 
@@ -1241,16 +1374,16 @@ def _nHint(self, z: npt.NDArray, d: npt.NDArray) -> npt.NDArray:
         """
         Internal method used to calcualte the basis function value.
 
-        Parameters:
-        -----------
-        z: NDArray
+        Parameters
+        ----------
+        z : NDArray
             Values to calculate the basis functions for.
-        d: NDArray
+        d : NDArray
             Derivative order.
 
-        Returns:
-        --------
-        H: NDArray
+        Returns
+        -------
+        H : NDArray
             Basis function values.
         """
         pass
@@ -1271,16 +1404,16 @@ def _nHint(self, z: npt.NDArray, d: npt.NDArray) -> npt.NDArray:
         """
         Internal method used to calcualte the basis function value.
 
-        Parameters:
-        -----------
-        z: NDArray
+        Parameters
+        ----------
+        z : NDArray
             Values to calculate the basis functions for.
-        d: NDArray
+        d : NDArray
             Derivative order.
 
-        Returns:
-        --------
-        H: NDArray
+        Returns
+        -------
+        H : NDArray
             Basis function values.
         """
         ind = -1
@@ -1314,16 +1447,16 @@ def _nHint(self, z: npt.NDArray, d: npt.NDArray) -> npt.NDArray:
         """
         Internal method used to calcualte the basis function value.
 
-        Parameters:
-        -----------
-        z: NDArray
+        Parameters
+        ----------
+        z : NDArray
             Values to calculate the basis functions for.
-        d: NDArray
+        d : NDArray
             Derivative order.
 
-        Returns:
-        --------
-        H: NDArray
+        Returns
+        -------
+        H : NDArray
             Basis function values.
         """
 
@@ -1333,7 +1466,9 @@ def _nHint(self, z: npt.NDArray, d: npt.NDArray) -> npt.NDArray:
 
         z = jnp.split(z, z.shape[1], axis=1)
 
-        def Recurse(dark, d, dim, dCurr=0):
+        def Recurse(
+            dark: Callable[[npt.NDArray], jnp.ndarray], d: uint, dim: uint, dCurr: uint = 0
+        ) -> Callable[[npt.NDArray], jnp.ndarray]:
             if dCurr == d:
                 return dark
             else:
@@ -1347,7 +1482,7 @@ def Recurse(dark, d, dim, dCurr=0):
             dark2 *= self._c[dim] ** deriv
             dark = Recurse(dark, deriv, dim)
 
-        return (dark(*z) * dark2).to_py()
+        return np.asarray((dark(*z) * dark2))
 
 
 class nELMTanh(nELM):
@@ -1359,15 +1494,15 @@ def _nHint(self, z: npt.NDArray, d: npt.NDArray) -> npt.NDArray:
         """
         Internal method used to calcualte the basis function value.
 
-        Parameters:
-        -----------
-        z: NDArray
+        Parameters
+        ----------
+        z : NDArray
             Values to calculate the basis functions for.
-        d: NDArray
+        d : NDArray
             Derivative order.
 
-        Returns:
-        --------
+        Returns
+        -------
         H: NDArray
             Basis function values.
         """
@@ -1378,7 +1513,9 @@ def _nHint(self, z: npt.NDArray, d: npt.NDArray) -> npt.NDArray:
 
         z = jnp.split(z, z.shape[1], axis=1)
 
-        def Recurse(dark, d, dim, dCurr=0):
+        def Recurse(
+            dark: Callable[[npt.NDArray], jnp.ndarray], d: uint, dim: uint, dCurr: uint = 0
+        ) -> Callable[[npt.NDArray], jnp.ndarray]:
             if dCurr == d:
                 return dark
             else:
@@ -1392,7 +1529,7 @@ def Recurse(dark, d, dim, dCurr=0):
             dark2 *= self._c[dim] ** deriv
             dark = Recurse(dark, deriv, dim)
 
-        return (dark(*z) * dark2).to_py()
+        return np.asarray((dark(*z) * dark2))
 
 
 class nELMSigmoid(nELM):
@@ -1423,7 +1560,9 @@ def _nHint(self, z: npt.NDArray, d: npt.NDArray) -> npt.NDArray:
 
         z = jnp.split(z, z.shape[1], axis=1)
 
-        def Recurse(dark, d, dim, dCurr=0):
+        def Recurse(
+            dark: Callable[[npt.NDArray], jnp.ndarray], d: uint, dim: uint, dCurr: uint = 0
+        ) -> Callable[[npt.NDArray], jnp.ndarray]:
             if dCurr == d:
                 return dark
             else:
@@ -1437,7 +1576,7 @@ def Recurse(dark, d, dim, dCurr=0):
             dark2 *= self._c[dim] ** deriv
             dark = Recurse(dark, deriv, dim)
 
-        return (dark(*z) * dark2).to_py()
+        return np.asarray((dark(*z) * dark2))
 
 
 class nELMSwish(nELM):
@@ -1449,16 +1588,16 @@ def _nHint(self, z: npt.NDArray, d: npt.NDArray) -> npt.NDArray:
         """
         Internal method used to calcualte the basis function value.
 
-        Parameters:
-        -----------
-        z: NDArray
+        Parameters
+        ----------
+        z : NDArray
             Values to calculate the basis functions for.
-        d: NDArray
+        d : NDArray
             Derivative order.
 
-        Returns:
-        --------
-        H: NDArray
+        Returns
+        -------
+        H : NDArray
             Basis function values.
         """
 
@@ -1470,7 +1609,9 @@ def f(*x):
 
         z = jnp.split(z, z.shape[1], axis=1)
 
-        def Recurse(dark, d, dim, dCurr=0):
+        def Recurse(
+            dark: Callable[[npt.NDArray], jnp.ndarray], d: uint, dim: uint, dCurr: uint = 0
+        ) -> Callable[[npt.NDArray], jnp.ndarray]:
             if dCurr == d:
                 return dark
             else:
@@ -1484,4 +1625,4 @@ def Recurse(dark, d, dim, dCurr=0):
             dark2 *= self._c[dim] ** deriv
             dark = Recurse(dark, deriv, dim)
 
-        return (dark(*z) * dark2).to_py()
+        return np.asarray((dark(*z) * dark2))
diff --git a/src/tfc/utils/Html.py b/src/tfc/utils/Html.py
index 6e99b4e..0d229d6 100644
--- a/src/tfc/utils/Html.py
+++ b/src/tfc/utils/Html.py
@@ -1,6 +1,8 @@
 import os
 from graphviz import Digraph
 from yattag import Doc, indent
+from .types import List
+from .types import Path
 
 
 class HTML:
@@ -13,12 +15,13 @@ class HTML:
         Output file
     """
 
-    def __init__(self, outFile):
-        """This function initializes the header file, and saves useful variables to self.
+    def __init__(self, outFile: Path):
+        """
+        This function initializes the header file, and saves useful variables to self.
 
         Parameters
         ----------
-        outFile: str
+        outFile : Path
             Output file
         """
         self._outFile = outFile
@@ -27,8 +30,16 @@ def __init__(self, outFile):
             ".center {\n\tdisplay: block;\n\tmargin-left: auto;\n\tmargin-right: auto;\n}\n"
         )
 
-    def GenerateHtml(self):
-        """This function generates and formats the HTML file text."""
+    def GenerateHtml(self) -> str:
+        """
+        This function generates and formats the HTML file text.
+
+        Returns:
+        --------
+        html : str
+            HTML file as a string.
+        """
+
         html = indent(self.doc.getvalue(), indentation="", newline="\n")
         return html
 
@@ -40,13 +51,13 @@ def WriteFile(self):
         out.write(self.GenerateHtml())
         out.close()
 
-    def ReadFile(self, inFile):
+    def ReadFile(self, inFile: Path) -> str:
         """This function reads the file specified by "inFile" and retuns the
         contents as a string.
 
         Parameters
         ----------
-        inFile : str
+        inFile : Path
             File to read.
 
         Returns
@@ -63,30 +74,33 @@ def ReadFile(self, inFile):
 class Dot:
     """
     This class contains helper functions used to create dot graphs.
+    """
 
-    Parameters
-    ----------
-    outFile : str
-        Name of the filename under which the dot file should be saved.
+    def __init__(self, outFile: Path, name: str):
+        """
+        This function initializes the class and creates the digraph.
 
-    name: str
-        What the dot file should be called by Digrapph.
+        Parameters
+        ----------
+        outFile : Path
+            Name of the filename under which the dot file should be saved.
 
-    """
+        name : str
+            What the dot file should be called by Digraph.
+        """
 
-    def __init__(self, outFile, name):
-        """This function initizes the class and creates the digraph."""
         self._outFile = outFile
         self._name = name
         self.dot = Digraph(name=self._name)
 
-    def Render(self, formats=["cmapx", "svg"]):
-        """This function renders the dot graph as a .svg and as a .cmapx.
+    def Render(self, formats: List[str] = ["cmapx", "svg"]):
+        """
+        This function renders the dot graph as a .svg and as a .cmapx.
 
         Parameters
         ----------
-        formats : list, optional
-            List whose elementts dictatte which formats to render the dot graph in. (Default value = ["cmapx", "svg"]
+        formats : List[str], optional
+            List whose elementts dictate which formats to render the dot graph in. (Default value = ["cmapx", "svg"])
         """
         for f in formats:
             self.dot.render(self._outFile, format=f, cleanup=True, view=False)
diff --git a/src/tfc/utils/Latex.py b/src/tfc/utils/Latex.py
index d682c6e..afa6041 100644
--- a/src/tfc/utils/Latex.py
+++ b/src/tfc/utils/Latex.py
@@ -1,11 +1,14 @@
 import numpy as np
+from numpy import typing as npt
+from .types import List
+from typing import Optional
 
 
 class table:
     """This class is used to create the text needed for latex tables."""
 
     @staticmethod
-    def _Header(numCols):
+    def _Header(numCols: int) -> str:
         """This function creates the table header based on the number of columns.
 
         Parameters
@@ -21,13 +24,13 @@ def _Header(numCols):
         return "\\begin{center}\n\\begin{tabular}{" + "|c" * numCols + "|}\n"
 
     @staticmethod
-    def _colHeader(strIn):
+    def _colHeader(strIn: List[str]) -> str:
         """This function creates the column header based on the list of strings that are passed in
         via the input strIn.
 
         Parameters
         ----------
-        strIn : list
+        strIn : List[str]
             List of strings that form the column headers.
 
         Returns
@@ -38,8 +41,9 @@ def _colHeader(strIn):
         return " & ".join(strIn) + "\\\\\n"
 
     @staticmethod
-    def _Arr2Tab(arrIn, form="%.4E", rowHeader=None):
-        """This function transforms the 2-D numpy array (arrIn) into latex tabular format. The
+    def _Arr2Tab(arrIn: npt.NDArray, form: str = "%.4E", rowHeader: Optional[List[str]] = None):
+        """
+        This function transforms the 2-D numpy array (arrIn) into latex tabular format. The
         "form" argument specifies the number format to be used in the tabular environment.
         The "rowHeader" argument is a list of strings that are used in the first column of
         each row in the tabular environment. The latex tabular environment is returned as a
@@ -47,13 +51,13 @@ def _Arr2Tab(arrIn, form="%.4E", rowHeader=None):
 
         Parameters
         ----------
-        arrIn : array-like
+        arrIn : NDArray
             Array to convert to Latex table.
 
         form : str, optional
              Format string for the table numbers. (Default value = "%.4E")
 
-        rowHeader : list, optional
+        rowHeader : Optional[List[str]]
              List of strings to use as the row headers. (Default value = None)
 
         Returns
@@ -92,12 +96,19 @@ def _Arr2Tab(arrIn, form="%.4E", rowHeader=None):
         return out.rstrip()
 
     @staticmethod
-    def _Footer():
-        """This function creates the footer for the latex table."""
+    def _Footer() -> str:
+        """
+        This function creates the footer for the latex table.
+        """
         return "\\end{tabular}\n\\end{center}"
 
     @staticmethod
-    def SimpleTable(arrIn, form="%.4E", colHeader=None, rowHeader=None):
+    def SimpleTable(
+        arrIn: npt.NDArray,
+        form: str = "%.4E",
+        colHeader: Optional[List[str]] = None,
+        rowHeader: Optional[List[str]] = None,
+    ) -> str:
         """This function creates a simple latex table for the 2D numpy array arrIn. The
         "form" argument specifies the number format to be used in the tabular environment.
         The "colHeader" arugment is a list of strings that are used as the first row in the
@@ -107,22 +118,22 @@ def SimpleTable(arrIn, form="%.4E", colHeader=None, rowHeader=None):
 
         Parameters
         ----------
-        arrIn : array-like
+        arrIn : NDArray
             Array to convert to Latex table.
 
         form : str, optional
              Format string for the table numbers. (Default value = "%.4E")
 
-        colHeader : list, optional
+        colHeader : Optional[List[str]]
             List of strings that form the column headers. (Default value = None)
 
-        rowHeader : list, optional
+        rowHeader : Optional[List[str]]
              List of strings to use as the row headers. (Default value = None)
 
         Returns
         -------
-            table : str
-                Latex table as a string.
+        table : str
+            Latex table as a string.
         """
 
         if colHeader is None and rowHeader is None:
diff --git a/src/tfc/utils/MakePlot.py b/src/tfc/utils/MakePlot.py
index b923873..51c77e7 100644
--- a/src/tfc/utils/MakePlot.py
+++ b/src/tfc/utils/MakePlot.py
@@ -3,38 +3,46 @@
 import numpy as np
 
 import matplotlib.pyplot as plt
-from mpl_toolkits.mplot3d import Axes3D
 
 from .TFCUtils import TFCPrint
+from .types import StrArrayLike, Path, List, Dict, Literal, pint
+from typing import Optional, Union, Generator, Callable
 
 TFCPrint()
 
 
 class MakePlot:
-    """This class is used to easily create journal-article-ready plots and subplots. The class can create 2D as well as 3D plots
+    """
+    This class is used to easily create journal-article-ready plots and subplots. The class can create 2D as well as 3D plots
     and even has support for twin y-axes.
-
-    Parameters
-    ----------
-    xlabs: list or array-like
-        The x-axes labels of for the plots
-    ylabs: list or array-like
-        The y-axes labels of for the plots
-    zlabs: list or array-like, optional
-        The z-axes labels of for the plots. Setting this forces subplots to be 3D. (Default value = None)
     """
 
-    def __init__(self, xlabs, ylabs, twinYlabs=None, titles=None, zlabs=None, style=None):
-        """This function initializes subplots based on the inputs provided.
+    def __init__(
+        self,
+        xlabs: StrArrayLike,
+        ylabs: StrArrayLike,
+        titles: Optional[StrArrayLike] = None,
+        twinYlabs: Optional[StrArrayLike] = None,
+        zlabs: Optional[StrArrayLike] = None,
+        style: Optional[Union[str, Dict, Path, List[str], List[Dict], List[Path]]] = None,
+    ):
+        """
+        This function initializes the plot/subplots based on the inputs provided.
 
         Parameters
         ----------
-        xlabs: list or array-like
-            The x-axes labels of for the plots
-        ylabs: list or array-like
-            The y-axes labels of for the plots
-        zlabs: list or array-like, optional
+        xlabs: StrArrayLike
+            The x-axes labels for the plots
+        ylabs: StrArrayLike
+            The y-axes labels for the plots
+        titles: StrArrayLike, optional
+            The titles for the plots. (Default value = None)
+        twinYlabs: StrArrayLike, optional
+            The twin y-axes labels for the plots. Setting this forces twin axis y-axes. (Default value = None)
+        zlabs: StrArrayLike, optional
             The z-axes labels of for the plots. Setting this forces subplots to be 3D. (Default value = None)
+        style : Union[str, Dict, Path, List[str], List[Dict], List[Path]]
+            Matplotlib style. (Default value = None)
         """
 
         # Apply a style if specified
@@ -179,7 +187,9 @@ def __init__(self, xlabs, ylabs, twinYlabs=None, titles=None, zlabs=None, style=
         self.fig.tight_layout()
 
     def FullScreen(self):
-        """This function makes the plot fullscreen."""
+        """
+        This function makes the plot fullscreen.
+        """
 
         # Get screensize
         import tkinter as tk
@@ -193,8 +203,9 @@ def FullScreen(self):
         dpi = float(self.fig.get_dpi())
         self.fig.set_size_inches(width / dpi, height / dpi)
 
-    def PartScreen(self, width, height):
-        """This function makes the plot width x height inches.
+    def PartScreen(self, width: float, height: float):
+        """
+        This function makes the plot width x height inches.
 
         Parameters
         ----------
@@ -207,25 +218,33 @@ def PartScreen(self, width, height):
         self.fig.set_size_inches(width, height)
 
     def show(self):
-        """This function shows the plot."""
+        """
+        This function shows the plot.
+        """
         self.fig.show()
 
     def draw(self):
-        """This function draws the canvas."""
+        """
+        This function draws the canvas.
+        """
         self.fig.canvas.draw()
 
-    def save(self, fileName, transparent=True, fileType="pdf"):
-        """This function crops and saves the figure.
+    def save(
+        self,
+        fileName: Path,
+        transparent: bool = True,
+        fileType: Literal["png", "pdf", "ps", "eps", "svg"] = "pdf",
+    ):
+        """
+        This function crops and saves the figure.
 
         Parameters
         ----------
-        fileName : str
+        fileName : Path
             Filename where the figure should be saved. Note, this should not include the file extension.
-
         transparent : bool, optional
             Whether to save the plot with transparency or not. (Default value = True)
-
-        fileType : str, optional
+        fileType : Literal["png", "pdf", "ps", "eps", "svg"], optional
             File exension to use. (Default value = "pdf")
         """
         self.fig.savefig(
@@ -237,51 +256,58 @@ def save(self, fileName, transparent=True, fileType="pdf"):
             transparent=transparent,
         )
 
-    def savePickle(self, fileName):
+    def savePickle(self, fileName: Path):
         """This function saves the figure in a pickle format so it can be opened and modified later.
 
         Parameters
         ----------
-        fileName : str
+        fileName : Path
             Filename where the figure should be saved. Note, this should not include the file extension.
         """
         pickle.dump(self.fig, open(fileName + ".pickle", "wb"))
 
-    def saveAll(self, fileName, transparent=True, fileType="pdf"):
+    def saveAll(
+        self,
+        fileName: Path,
+        transparent: bool = True,
+        fileType: Literal["png", "pdf", "ps", "eps", "svg"] = "pdf",
+    ):
         """This function invokes the save and savePickle functions.
 
         Parameters
         ----------
-        fileName : str
+        fileName : Path
             Filename where the figure should be saved. Note, this should not include the file extension.
-
         transparent : bool, optional
             Whether to save the plot with transparency or not. (Default value = True)
-
-        fileType : str, optional
+        fileType : Literal["png", "pdf", "ps", "eps", "svg"], optional
             File exension to use. (Default value = "pdf")
         """
         self.save(fileName, transparent=transparent, fileType=fileType)
         self.savePickle(fileName)
 
-    def animate(self, animFunc, outDir="MyMovie", fileName="images", save=True, delay=10):
-        """Creates an animation using a Python generator.
+    def animate(
+        self,
+        animFunc: Callable[[], Generator[None, None, None]],
+        outDir: Path = "MyMovie",
+        fileName: str = "images",
+        save: bool = True,
+        delay: pint = 10,
+    ):
+        """
+        Creates an animation using a Python generator.
 
         Parameters
         ----------
-        animFunc : generator function
-            Function that modifies what is displayed on the plot
-
-        outDir : str, optional
+        animFunc: Callable[[], Generator[None, None, None]]
+            Function whose returned generator modifies what is displayed on the plot
+        outDir : Path, optional
              Directory to save frames in: only used if save = True. (Default value = "MyMovie")
-
         fileName : str, optional
              Filename for the frames: only used if save = True. (Default value = "images")
-
         save : bool, optional
              Controls whether the function saves the frames of the animation or not. (Default value = True)
-
-        delay : integer, optional
+        delay : pint, optional
              Delay in milliseconds between frames: only used if save = False. (Default value = 10)
         """
 
diff --git a/src/tfc/utils/MayaviMakePlot.py b/src/tfc/utils/MayaviMakePlot.py
index b58d671..f52d16e 100644
--- a/src/tfc/utils/MayaviMakePlot.py
+++ b/src/tfc/utils/MayaviMakePlot.py
@@ -1,13 +1,19 @@
 import numpy as np
+import numpy.typing as npt
+import mayavi
 from mayavi import mlab
 from matplotlib import colors as mcolors
+from .types import Dict, Tuple, Path, Ge, Le, Annotated, Literal
+from typing import Optional, Any, Union, Generator, Callable
+
+Color = Union[str, Tuple[float, float, float, float], npt.NDArray[np.float64]]
 
 
 class MakePlot:
     """MakePlot class for Mayavi."""
 
     @staticmethod
-    def _str_to_rgb(color):
+    def _str_to_rgb(color: str) -> Tuple[float, float, float]:
         """Call matplotlib's colorConverter.to_rgb on input string.
 
         Parameters
@@ -15,16 +21,15 @@ def _str_to_rgb(color):
         color : str
             Color string
 
-
         Returns
         -------
-        color_rgb : tuple
+        color_rgb : Tuple[float, float, float]
             3-tuple of the RGB for the color
         """
         return mcolors.colorConverter.to_rgb(color)
 
     @staticmethod
-    def _str_to_rgba(color, alpha=None):
+    def _str_to_rgba(color, alpha: Optional[float] = None) -> Tuple[float, float, float, float]:
         """Call matplotlib's colorConverter.to_rgba on input string.
 
         Parameters
@@ -37,23 +42,23 @@ def _str_to_rgba(color, alpha=None):
 
         Returns
         -------
-        color_rgba : tuple
+        color_rgba : Tuple[float, float, float, float]
             4-tuple of the RGB for the color
         """
-        return mcolors.colorConverter.to_rgba(color, alpha=None)
+        return mcolors.colorConverter.to_rgba(color, alpha=alpha)
 
     @staticmethod
-    def _ProcessKwargs(**kwargs):
+    def _ProcessKwargs(**kwargs: Any) -> Dict[str, Any]:
         """This function effectively extends common mlab keywords.
 
         Parameters
         ----------
-        **kwargs : dict
+        **kwargs : Any
             keyword arguments
 
         Returns
         -------
-        kwargs : dict
+        kwargs : Dict[str, any]
             Same as input keyword arguments but color has been transformed to an RGB if it was a string.
         """
         # Process color argument if it exists
@@ -64,54 +69,62 @@ def _ProcessKwargs(**kwargs):
         return kwargs
 
     @staticmethod
-    def ColormapGradient(c1, c2):
+    def ColormapGradient(c1: Color, c2: Color) -> npt.NDArray[np.uint8]:
         """Returns a Mayavi LUT for a gradient that linearly transforms from c1 to c2.
 
         Parameters
         ----------
-        c1 : str or list
+        c1 : Color
             Input color as a string or 4-element RGBA list.
 
-        c2 : str or list
+        c2 : Color
             Input color as a string or 4-element RGBA list.
 
         Returns
         -------
-        lut : array-like
+        lut : npt.NDArray[np.uint8]
             Mayavi LUT for a gradient that linearly transforms from c1 to c2.
-
         """
         if isinstance(c1, str):
             c1 = np.array(MakePlot._str_to_rgba(c1))
         if isinstance(c2, str):
             c2 = np.array(MakePlot._str_to_rgba(c2))
+        if not isinstance(c1, np.ndarray):
+            c1 = np.array(c1)
+        if not isinstance(c2, np.ndarray):
+            c2 = np.array(c2)
         c1 = c1 * 255
         c2 = c2 * 255
         return np.linspace(c1, c2, 256, dtype=np.uint8)
 
     @staticmethod
-    def SetColormapGradient(obj, c1, c2):
-        """Applies a linear, gradient colormap to the object. The colors in the gradient are c1 and c2.
+    def SetColormapGradient(obj: Any, c1: Color, c2: Color):
+        """
+        Applies a linear, gradient colormap to the object. The colors in the gradient are c1 and c2.
 
         Parameters
         ----------
-        obj : mlab object
+        obj : Any
             mlab object to apply the colormap gradient to.
-
-        c1 : str or list
+        c1 : Color
             Input color as a string or 4-element RGBA list.
-
-        c2 : str or list
+        c2 : Color
             Input color as a string or 4-element RGBA list.
         """
         obj.module_manager.scalar_lut_manager.lut.table = MakePlot.ColormapGradient(c1, c2)
 
     def __init__(self):
+        """
+        This function initializes the plot.
+        """
+
         self.fig = mlab.figure(bgcolor=(1.0, 1.0, 1.0))
         self.scene = self.fig.scene
 
     def FullScreen(self):
-        """This function makes the plot fullscreen."""
+        """
+        This function makes the plot fullscreen.
+        """
 
         # Get screensize
         import tkinter as tk
@@ -124,7 +137,7 @@ def FullScreen(self):
         dpi = float(self.fig.get_dpi())
         mlab.figure(self.fig, size=(width / dpi, height / dpi))
 
-    def PartScreen(self, width, height):
+    def PartScreen(self, width: float, height: float):
         """This function makes the plot width x height inches.
 
         Parameters
@@ -137,29 +150,34 @@ def PartScreen(self, width, height):
         """
         mlab.figure(self.fig, size=(width, height))
 
-    def animate(self, animFunc, outDir="MyMovie", fileName="images", save=True, delay=10):
+    IntRange = Annotated[int, Ge(10), Le(100000)]
+
+    def animate(
+        self,
+        animFunc: Callable[[], Generator[None, None, None]],
+        outDir: Path = "MyMovie",
+        fileName: str = "images",
+        save: bool = True,
+        delay: IntRange = 10,
+    ) -> Optional[mayavi.tools.animator.Animator]:
         """
 
         Parameters
         ----------
-        animFunc : function generator
+        animFunc : Callable[[], Generator[None, None, None]]
             Function that modifies what is displayed on the plot
-
-        outDir : str, optional
-             Directory to save frames in: only used if save = True. (Default value = "MyMovie")
-
+        outDir : Path, optional
+            Directory to save frames in: only used if save = True. (Default value = "MyMovie")
         fileName : str, optional
-             Filename for the frames: only used if save = True. (Default value = "images")
-
+            Filename for the frames: only used if save = True. (Default value = "images")
         save : bool, optional
-             Controls whether the function saves the frames of the animation or not. (Default value = True)
-
-        delay: int in [10,100000], optional
-             Controls the delay used by mlab.animate. (Default value = 10)
+            Controls whether the function saves the frames of the animation or not. (Default value = True)
+        delay: IntRange in [10,100000], optional
+            Controls the delay used by mlab.animate. (Default value = 10)
 
         Returns
         -------
-        animator : mayavi.tools.animator.Animator
+        animator : Optional[mayavi.tools.animator.Animator]
             Returns an animator object only if save = False. Otherwise, there is no return value.
         """
         if save:
@@ -183,15 +201,35 @@ def animate(self, animFunc, outDir="MyMovie", fileName="images", save=True, dela
         else:
             return mlab.animate(func=animFunc, delay=delay)()
 
-    def save(self, fileName, fileType="pdf"):
-        """This function saves the figure.
+    def save(
+        self,
+        fileName: Path,
+        fileType: Literal[
+            "png",
+            "jpg",
+            "bmp",
+            "tiff",
+            "ps",
+            "eps",
+            "pdf",
+            "rib",
+            "oogl",
+            "iv",
+            "wrl",
+            "vrml",
+            "obj",
+            "x3d",
+            "pov",
+        ] = "pdf",
+    ):
+        """
+        This function saves the figure.
 
         Parameters
         ----------
-        fileName : str
+        fileName : Path
             Filename where the figure should be saved. Note, this should not include the file extension.
-
-        fileType : str, optional
+        fileType : Literal["png", "jpg", "bmp", "tiff", "ps", "eps", "pdf", "rib", "oogl", "iv", "wrl", "vrml", "obj", "x3d", "pov"], optional
             File exension to use. (Default value = "pdf")
         """
         mlab.savefig(fileName + "." + fileType, figure=self.fig)
@@ -202,113 +240,143 @@ def show(self):
 
     @property
     def show_axes(self):
-        """Axes indicator."""
+        """
+        Adds axes indicator.
+        """
         return self.scene.show_axes
 
     @show_axes.setter
-    def show_axes(self, val):
-        """Set axes indicator.
+    def show_axes(self, val: bool):
+        """
+        Set axes indicator.
 
         Parameters
         ----------
         val : bool
-            True turn on the axes indicator, and False turns off the axes indicator.
+            True turns on the axes indicator and False turns off the axes indicator.
         """
         self.scene.show_axes = val
 
-    def view(self, *args, **kwargs):
-        """Call mlab's view on the class's figure.
+    def view(self, *args: Any, **kwargs: Any):
+        """
+        Call mlab's view on the class's figure.
 
         Parameters
         ----------
-        *args : iterable
+        *args : Any
             args passed on to view
-
-        **kwargs : dict
+        **kwargs : Any
             kwargs passed on to view
         """
         return mlab.view(*args, figure=self.fig, **kwargs)
 
-    def points3d(self, *args, **kwargs):
-        """Call mlab's points3d on the class's figure.
+    def points3d(self, *args: Any, **kwargs: Any) -> mayavi.modules.glyph.Glyph:
+        """
+        Call mlab's points3d on the class's figure.
 
         Parameters
         ----------
-        *args : iterable
+        *args : Any
             args passed on to points3d
-
-        **kwargs : dict
+        **kwargs : Any
             kwargs passed on to points3d after being processed by _ProcessKwargs.
+
+        Returns:
+        --------
+        points3d : mayavi.modules.glyph.Glyph
+            Glyphs corresponding to the points.
         """
         kwargs = MakePlot._ProcessKwargs(**kwargs)
         return mlab.points3d(*args, figure=self.fig, **kwargs)
 
-    def plot3d(self, *args, **kwargs):
+    def plot3d(self, *args: Any, **kwargs: Any) -> mayavi.modules.surface.Surface:
         """Call mlab's plot3d on the class's figure.
 
         Parameters
         ----------
-        *args : iterable
+        *args : Any
             args passed on to plot3d
-
-        **kwargs : dict
+        **kwargs : Any
             kwargs passed on to plot3d after being processed by _ProcessKwargs.
+
+        Returns:
+        --------
+        surf : mayavi.modules.surface.Surface
+            Line as a Mayavi surface.
         """
         kwargs = MakePlot._ProcessKwargs(**kwargs)
         return mlab.plot3d(*args, figure=self.fig, **kwargs)
 
-    def surf(self, *args, **kwargs):
+    def surf(self, *args: Any, **kwargs: Any) -> mayavi.modules.surface.Surface:
         """Call mlab's surf on the class's figure.
 
         Parameters
         ----------
-        *args : iterable
+        *args : Any
             args passed on to surf
-
-        **kwargs : dict
+        **kwargs : Any
             kwargs passed on to surf after being processed by _ProcessKwargs.
+
+        Returns:
+        --------
+        surf : mayavi.modules.surface.Surface
+            Surface.
         """
         kwargs = MakePlot._ProcessKwargs(**kwargs)
         return mlab.surf(*args, figure=self.fig, **kwargs)
 
-    def quiver3d(self, *args, **kwargs):
-        """Call mlab's quiver3d on the class's figure.
+    def quiver3d(self, *args: Any, **kwargs: Any) -> mayavi.modules.vectors.Vectors:
+        """
+        Call mlab's quiver3d on the class's figure.
 
         Parameters
         ----------
-        *args : iterable
+        *args : Any
             args passed on to quiver3d
-
-        **kwargs : dict
+        **kwargs : Any
             kwargs passed on to quiver3d after being processed by _ProcessKwargs.
+
+        Returns:
+        --------
+        vectors : mayavi.modules.vectors.Vectors
+            Vectors associated with the quiver3d plot.
         """
         kwargs = MakePlot._ProcessKwargs(**kwargs)
         return mlab.quiver3d(*args, figure=self.fig, **kwargs)
 
-    def mesh(self, *args, **kwargs):
+    def mesh(self, *args: Any, **kwargs: Any) -> mayavi.modules.surface.Surface:
         """Call mlab's mesh on the class's figure.
 
         Parameters
         ----------
-        *args : iterable
+        *args : Any
             args passed on to mesh
-
-        **kwargs : dict
+        **kwargs : Any
             kwargs passed on to mesh after being processed by _ProcessKwargs.
+
+        Returns:
+        --------
+        mesh : mayavi.modules.surface.Surface
+            Mesh as a mayavi surface.
         """
         kwargs = MakePlot._ProcessKwargs(**kwargs)
         return mlab.mesh(*args, figure=self.fig, **kwargs)
 
-    def contour3d(self, *args, **kwargs):
-        """Call mlab's contour3d on the class's figure.
+    def contour3d(self, *args: Any, **kwargs: Any) -> mayavi.modules.iso_surface.IsoSurface:
+        """
+        Call mlab's contour3d on the class's figure.
 
         Parameters
         ----------
-        *args : iterable
+        *args : Any
             args passed on to contour3d
-
-        **kwargs : dict
+        **kwargs : Any
             kwargs passed on to mesh after being processed by _ProcessKwargs.
+
+        Returns:
+        --------
+        contour : mayavi.modules.iso_surface.IsoSurface
+            Coutour.
         """
         kwargs = MakePlot._ProcessKwargs(**kwargs)
         return mlab.contour3d(*args, figure=self.fig, **kwargs)
diff --git a/src/tfc/utils/PlotlyMakePlot.py b/src/tfc/utils/PlotlyMakePlot.py
index 3ff3b95..d23707a 100644
--- a/src/tfc/utils/PlotlyMakePlot.py
+++ b/src/tfc/utils/PlotlyMakePlot.py
@@ -1,27 +1,16 @@
 import numpy as np
+import plotly
 import plotly.graph_objects as go
 
 from .TFCUtils import TFCPrint
+from .types import StrArrayLike, uint, Path, Literal, List
+from typing import Optional
 
 
 class MakePlot:
     """
     A MakePlot class for Plotly.
     This class wraps common Plotly functions to ease figure creation.
-
-    Parameters
-    ----------
-    xlabs: list or array-like
-        The x-axes labels of for the plots
-
-    ylabs: list or array-like
-        The y-axes labels of for the plots
-
-    zlabs: list or array-like, optional
-        The z-axes labels of for the plots. Setting this forces subplots to be 3D. (Default value = None)
-
-    titles: list or array-like, optional
-        The titles for the plots. (Default value = None)
     """
 
     _template = {
@@ -87,7 +76,27 @@ class MakePlot:
     _backgroundColor = "rgba(0,0,0,0)"
     _gridColor = "rgb(176, 176, 176)"
 
-    def __init__(self, xlabs, ylabs, titles=None, zlabs=None):
+    def __init__(
+        self,
+        xlabs: StrArrayLike,
+        ylabs: StrArrayLike,
+        titles: Optional[StrArrayLike] = None,
+        zlabs: Optional[StrArrayLike] = None,
+    ):
+        """
+        This function initializes the plot/subplots based on the inputs provided.
+
+        Parameters
+        ----------
+        xlabs: StrArrayLike
+            The x-axes labels for the plots
+        ylabs: StrArrayLike
+            The y-axes labels for the plots
+        titles: StrArrayLike, optional
+            The titles for the plots. (Default value = None)
+        zlabs: StrArrayLike, optional
+            The z-axes labels of for the plots. Setting this forces subplots to be 3D. (Default value = None)
+        """
 
         # View distance used by the view method
         self.viewDistance = 2.5
@@ -280,59 +289,65 @@ def __init__(self, xlabs, ylabs, titles=None, zlabs=None):
             template=self._template,
         )
 
-    def Surface(self, row=None, col=None, **kwargs):
+    def Surface(
+        self, row: Optional[uint] = None, col: Optional[uint] = None, **kwargs
+    ) -> plotly.graph_objects.Surface:
         """
         Creates a plotly surface on the subplot specified by row and col or on
         the main figure if not using subplots.
 
         Parameters
         ----------
-        row : int, optional
-             subplot row (Default value = None)
-        col : int, optional
-             subplot column (Default value = None)
+        row : Optional[uint]
+            Subplot row (Default value = None)
+        col : Optional[uint]
+            Subplot column (Default value = None)
         **kwargs : dict, optional
-            keyword arguments passed on to plotly.graphic_objects.Surface
+            Keyword arguments passed on to plotly.graphic_objects.Surface
 
         Returns
         -------
-        surf : plotly.graphic_objects.Surface
+        surf : plotly.graph_objects.Surface
         """
         return self.fig.add_trace(go.Surface(**kwargs), row=row, col=col)
 
-    def Scatter3d(self, row=None, col=None, **kwargs):
+    def Scatter3d(
+        self, row: Optional[uint] = None, col: Optional[uint] = None, **kwargs
+    ) -> plotly.graph_objects.Scatter3d:
         """
         Creates a 3d plotly scatter on the subplot specified by row and col or on
         the main figure if not using subplots.
 
         Parameters
         ----------
-        row : int, optional
-             subplot row (Default value = None)
-        col : int, optional
-             subplot column (Default value = None)
+        row : Optional[uint]
+             Subplot row (Default value = None)
+        col : Optional[uint]
+             Subplot column (Default value = None)
         **kwargs : dict, optional
             keyword arguments passed on to plotly.graphic_objects.Scatter3d
 
         Returns
         -------
-        scatter : plotly.graphic_objects.Scatter3d
+        scatter : plotly.graph_objects.Scatter3d
         """
         return self.fig.add_trace(go.Scatter3d(**kwargs), row=row, col=col)
 
-    def Scatter(self, row=None, col=None, **kwargs):
+    def Scatter(
+        self, row: Optional[uint] = None, col: Optional[uint] = None, **kwargs
+    ) -> plotly.graph_objects.Scatter:
         """
         Creates a plotly scatter on the subplot specified by row and col or on
         the main figure if not using subplots.
 
         Parameters
         ----------
-        row : int, optional
-             subplot row (Default value = None)
-        col : int, optional
-             subplot column (Default value = None)
+        row : Optional[uint]
+             Subplot row (Default value = None)
+        col : Optional[uint]
+             Subplot column (Default value = None)
         **kwargs : dict, optional
-            keyword arguments passed on to plotly.graphic_objects.Scatter
+            keyword arguments passed on to plotly.graph_objects.Scatter
 
         Returns
         -------
@@ -340,19 +355,21 @@ def Scatter(self, row=None, col=None, **kwargs):
         """
         return self.fig.add_trace(go.Scatter(**kwargs), row=row, col=col)
 
-    def Histogram(self, row=None, col=None, **kwargs):
+    def Histogram(
+        self, row: Optional[uint] = None, col: Optional[uint] = None, **kwargs
+    ) -> plotly.graph_objects.Histogram:
         """
         Creates a plotly histogram on the subplot specified by row and col or on
         the main figure if not using subplots.
 
         Parameters
         ----------
-        row : int, optional
-             subplot row (Default value = None)
-        col : int, optional
-             subplot column (Default value = None)
+        row : Optional[uint]
+             Subplot row (Default value = None)
+        col : Optional[uint]
+             Subplot column (Default value = None)
         **kwargs : dict, optional
-            keyword arguments passed on to plotly.graphic_objects.Histogram
+            keyword arguments passed on to plotly.graph_objects.Histogram
 
         Returns
         -------
@@ -360,39 +377,43 @@ def Histogram(self, row=None, col=None, **kwargs):
         """
         return self.fig.add_trace(go.Histogram(**kwargs), row=row, col=col)
 
-    def Contour(self, row=None, col=None, **kwargs):
+    def Contour(
+        self, row: Optional[uint] = None, col: Optional[uint] = None, **kwargs
+    ) -> plotly.graph_objects.Contour:
         """
         Creates a plotly contour on the subplot specified by row and col or on
         the main figure if not using subplots.
 
         Parameters
         ----------
-        row : int, optional
-             subplot row (Default value = None)
-        col : int, optional
-             subplot column (Default value = None)
+        row : Optional[uint]
+             Subplot row (Default value = None)
+        col : Optional[uint]
+             Subplot column (Default value = None)
         **kwargs : dict, optional
             keyword arguments passed on to plotly.graphic_objects.Contour
 
         Returns
         -------
-        contour : plotly.graphic_objects.Contour
+        contour : plotly.graph_objects.Contour
         """
         return self.fig.add_trace(go.Contour(**kwargs), row=row, col=col)
 
-    def Box(self, row=None, col=None, **kwargs):
+    def Box(
+        self, row: Optional[uint] = None, col: Optional[uint] = None, **kwargs
+    ) -> plotly.graph_objects.Box:
         """
         Creates a plotly box on the subplot specified by row and col or on
         the main figure if not using subplots.
 
         Parameters
         ----------
-        row : int, optional
-             subplot row (Default value = None)
-        col : int, optional
-             subplot column (Default value = None)
+        row : Optional[uint]
+             Subplot row (Default value = None)
+        col : Optional[uint]
+             Subplot column (Default value = None)
         **kwargs : dict, optional
-            keyword arguments passed on to plotly.graphic_objects.Box
+            keyword arguments passed on to plotly.graph_objects.Box
 
         Returns
         -------
@@ -400,19 +421,21 @@ def Box(self, row=None, col=None, **kwargs):
         """
         return self.fig.add_trace(go.Box(**kwargs), row=row, col=col)
 
-    def Violin(self, row=None, col=None, **kwargs):
+    def Violin(
+        self, row: Optional[uint] = None, col: Optional[uint] = None, **kwargs
+    ) -> plotly.graph_objects.Violin:
         """
         Creates a plotly violin on the subplot specified by row and col or on
         the main figure if not using subplots.
 
         Parameters
         ----------
-        row : int
-             subplot row (Default value = None)
-        col : int
-             subplot column (Default value = None)
+        row : Optional[uint]
+             Subplot row (Default value = None)
+        col : Optional[uint]
+             Subplot column (Default value = None)
         **kwargs : dict, optional
-            keyword arguments passed on to plotly.graphic_objects.Violin
+            keyword arguments passed on to plotly.graph_objects.Violin
 
         Returns
         -------
@@ -420,19 +443,21 @@ def Violin(self, row=None, col=None, **kwargs):
         """
         return self.fig.add_trace(go.Violin(**kwargs), row=row, col=col)
 
-    def Volume(self, row=None, col=None, **kwargs):
+    def Volume(
+        self, row: Optional[uint] = None, col: Optional[uint] = None, **kwargs
+    ) -> plotly.graph_objects.Volume:
         """
         Creates a plotly volume on the subplot specified by row and col or on
         the main figure if not using subplots.
 
         Parameters
         ----------
-        row : int
-             subplot row (Default value = None)
-        col : int
-             subplot column (Default value = None)
+        row : Optional[uint]
+             Subplot row (Default value = None)
+        col : Optional[uint]
+             Subplot column (Default value = None)
         **kwargs : dict, optional
-            keyword arguments passed on to plotly.graphic_objects.Volume
+            keyword arguments passed on to plotly.graph_objects.Volume
 
         Returns
         -------
@@ -442,7 +467,7 @@ def Volume(self, row=None, col=None, **kwargs):
 
     def show(self, **kwargs):
         """
-        Calls the figures show method.
+        Calls the figure's show method.
 
         Parameters
         ----------
@@ -451,20 +476,24 @@ def show(self, **kwargs):
         """
         self.fig.show(**kwargs)
 
-    def save(self, fileName, tight=True, fileType="png", **kwargs):
+    def save(
+        self,
+        fileName: Path,
+        tight: bool = True,
+        fileType: Literal["pdf", "jpg", "png", "svg", "eps", "html"] = "png",
+        **kwargs,
+    ):
         """
         Saves the figure using the type specified. If HTML is specified, the figure will
         be saved as a dynamic html file. All other file types are static.
 
         Parameters
         ----------
-        fileName : str
+        fileName : Path
             Name of the save file minus the file type (e.g., MyFigure not MyFigure.pdf).
-
         tight : boolean, optional
-            If the fileType is pdf or png and this value is true, then a tool is used to eliminate whitespace.
-            pdfCropMargins is used for PDFs and PIL is used for png's. (Default value = True)
-        fileType : {"pdf","jpg","png","svg","eps","html"}, optional
+            If the fileType is pdf or png and this value is true, then a tool is used to eliminate whitespace. pdfCropMargins is used for PDFs and PIL is used for png's. (Default value = True)
+        fileType : Literal["pdf","jpg","png","svg","eps","html"], optional
             Type of file to save the figure as. (Default value = "png")
         **kwargs : dict, optional
             Keyword arguments passed onto fig.write_image or fig.write_html, depending on fileType.
@@ -496,7 +525,7 @@ def save(self, fileName, tight=True, fileType="png", **kwargs):
 
                 crop(["-p", "0", fileNameFull])
 
-    def UploadToPlotly(self, username, apiKey, fileName, autoOpen=False):
+    def UploadToPlotly(self, username: str, apiKey: str, fileName: str, autoOpen: bool = False):
         """
         Upload your plot to Plotly.
 
@@ -504,13 +533,10 @@ def UploadToPlotly(self, username, apiKey, fileName, autoOpen=False):
         ----------
         username : str
             Plotly username
-
         apiKey : str
             Plotly api_key
-
         fileName : str
             Name of the file to save the plot as.
-
         autoOpen : bool, optional
             If true, plot will open in browser after saving. (Default value = False)
         """
@@ -520,7 +546,14 @@ def UploadToPlotly(self, username, apiKey, fileName, autoOpen=False):
         chart_studio.tools.set_credentials_file(username=username, api_key=apiKey)
         return chart_studio.plotly.plot(self.fig, filename=fileName, auto_open=autoOpen)
 
-    def view(self, azimuth, elevation, row=None, col=None, viewDistance=None):
+    def view(
+        self,
+        azimuth: float,
+        elevation: float,
+        row: Optional[uint] = None,
+        col: Optional[uint] = None,
+        viewDistance: Optional[float] = None,
+    ):
         """
         Change the view on the subplot specified by row and col or on
         the main figure if not using subplots.
@@ -531,12 +564,11 @@ def view(self, azimuth, elevation, row=None, col=None, viewDistance=None):
             Azimuth value in degrees
         elevation : float
             Elevation value in degrees
-
-        row : int, optional
-            subplot row (Default value = None)
-        col : int, optional
-            subplot column (Default value = None)
-        viewDistance : float, optional
+        row : Optional[uint]
+             Subplot row (Default value = None)
+        col : Optional[uint]
+             Subplot column (Default value = None)
+        viewDistance : Optional[float]
             Distance from camera to plot. (Default value = self.viewDistance)
         """
 
@@ -569,7 +601,9 @@ def view(self, azimuth, elevation, row=None, col=None, viewDistance=None):
             self.fig["layout"]["scene"]["camera"].eye = dict(x=dark[0], y=dark[1], z=dark[2])
 
     def FullScreen(self):
-        """Make the plot full screen."""
+        """
+        Make the plot full screen.
+        """
         import tkinter as tk
 
         root = tk.Tk()
@@ -579,7 +613,7 @@ def FullScreen(self):
 
         self.fig.update_layout(width=width, height=height)
 
-    def PartScreen(self, width, height, units="in"):
+    def PartScreen(self, width: float, height: float, units: Literal["in", "mm", "px"] = "in"):
         """
         Make the plot size equal to width x height.
 
@@ -589,8 +623,7 @@ def PartScreen(self, width, height, units="in"):
             Width of the plot
         height : float
             Height of the plot
-
-        units : {"in","mm","px"}, optional
+        units : Literal["in","mm","px"], optional
             Units width and height are given in. (Default value = inches)
         """
         if units != "px":
@@ -612,7 +645,13 @@ def PartScreen(self, width, height, units="in"):
         else:
             self.fig.update_layout(width=width, height=height)
 
-    def NormalizeColorScale(self, types=[], data=None, cmax=None, cmin=None):
+    def NormalizeColorScale(
+        self,
+        types: List[str] = [],
+        data: Optional[str] = None,
+        cmax: Optional[float] = None,
+        cmin: Optional[float] = None,
+    ):
         """
         Normalizes the color scale for the plots whose type is in the types list.
 
@@ -625,15 +664,13 @@ def NormalizeColorScale(self, types=[], data=None, cmax=None, cmin=None):
 
         Parameters
         ----------
-        types: list of str
+        types: List[str]
             Plot types to set cmax and cmin for.
-
-        data: str, optional
-            Data type to use to calculate cmax and cmin if not already specified.
-            (Default value = None)
-        cmax: float, optional
+        data: Optional[str]
+            Data type to use to calculate cmax and cmin if not already specified. (Default value = None)
+        cmax: Optional[float]
             cmax value to use when setting the colorscale
-        cmin: float, optional
+        cmin: Optional[float]
             cmin value to use when setting the colorscale
         """
 
diff --git a/src/tfc/utils/TFCUtils.py b/src/tfc/utils/TFCUtils.py
index 00fff66..d8cd4db 100644
--- a/src/tfc/utils/TFCUtils.py
+++ b/src/tfc/utils/TFCUtils.py
@@ -6,32 +6,48 @@
 from collections import OrderedDict
 from functools import partial
 
-from jax.config import config
+from jax._src.config import config
 
 config.update("jax_enable_x64", True)
 import numpy as onp
+import numpy.typing as npt
 import jax.numpy as np
 from jax import jvp, jit, lax, jacfwd, tree_map
 from jax import linear_util as lu
 from jax.util import safe_zip
 from jax.tree_util import register_pytree_node, tree_map
-from jax._src.api_util import flatten_fun, tree_flatten
+from jax._src.api_util import flatten_fun
+from jax._src.tree_util import tree_flatten
 from jax.core import get_aval, eval_jaxpr
-from jax.interpreters.partial_eval import JaxprTracer, trace_to_jaxpr_nounits, PartialVal
+from jax.interpreters.partial_eval import trace_to_jaxpr_nounits, PartialVal
 from jax.experimental.host_callback import id_tap
-from typing import List, Any, Callable
+from typing import Any, Callable, Optional, Union, cast
+from .types import uint, List, Literal, Tuple, TypedDict, Path, Dict
+from jaxtyping.pytree_type import PyTree
+from typing import cast
+
+# Types that can be added to a TFCDict
+TFCDictAddable = Union[np.ndarray, Dict[Any, Any], "TFCDict"]
+
+# Types that can be added to a TFCDictRobust
+TFCDictRobustAddable = Union[np.ndarray, Dict[Any, Any], "TFCDictRobust"]
 
 
 class TFCPrint:
-    """This class is used to print to the terminal in color."""
+    """
+    This class is used to print to the terminal in color.
+    """
 
     def __init__(self):
-        """This function is the constructor. It initializes the colorama class."""
+        """
+        This function is the constructor. It initializes the colorama class.
+        """
         initColorama()
 
     @staticmethod
-    def Error(stringIn):
-        """This function prints errors. It prints the text in 'stringIn' in bright red and
+    def Error(stringIn: str):
+        """
+        This function prints errors. It prints the text in 'stringIn' in bright red and
         exits the program.
 
         Parameters
@@ -44,8 +60,9 @@ def Error(stringIn):
         sys.exit()
 
     @staticmethod
-    def Warning(stringIn):
-        """This function prints warnings. It prints the text in 'stringIn' in bright yellow.
+    def Warning(stringIn: str):
+        """
+        This function prints warnings. It prints the text in 'stringIn' in bright yellow.
 
         Parameters
         ----------
@@ -56,15 +73,16 @@ def Warning(stringIn):
         print(style.RESET_ALL, end="")
 
 
-def egrad(g, j=0):
-    """This function mimics egrad from autograd.
+def egrad(g: Callable[..., Any], j: uint = 0):
+    """
+    This function mimics egrad from autograd.
 
     Parameters
     ----------
-    g : function
+    g : Callable[..., Any]
         Function to take the derivative of.
 
-    j : integer, optional
+    j : uint, optional
         Parameter with which to take the derivative with respect to. (Default value = 0)
 
     Returns
@@ -73,18 +91,18 @@ def egrad(g, j=0):
         Derivative function
     """
 
-    def wrapped(*args):
+    def wrapped(*args: Any) -> Any:
         """
         Wrapper for derivative of g with respect to parameter number j.
 
         Parameters
         ----------
-        *args : iterable
+        *args : Any
             function arguments to g
 
         Returns
         -------
-        x_bar: array-like
+        x_bar: Any
             derivative of g with respect to parameter number j
         """
         tans = tuple(
@@ -100,39 +118,51 @@ def wrapped(*args):
 
 
 @partial(partial, tree_map)
-def onesRobust(val):
-    """Returns ones_like val, but can handle arrays and dictionaries.
+def onesRobust(val: PyTree):
+    """
+    Returns ones_like val, but can handle arrays and dictionaries.
 
     Parameters
     ----------
-    val : pytree
+    val : PyTree
 
     Returns
     -------
-    ones_like_val : pytree
+    ones_like_val : PyTree
         Pytree with the same structure as val with all elements equal to one.
 
     """
+
+    # The @partial will force the PyTree to look array-like so we can use array
+    # functions as done below. Using an explicit cast here so LSPs will not complain.
+    val = cast(npt.NDArray, val)
+
     return onp.ones(val.shape, dtype=val.dtype)
 
 
 @partial(partial, tree_map)
-def zerosRobust(val):
-    """Returns zeros_like val, but can handle arrays and dictionaries.
+def zerosRobust(val: PyTree):
+    """
+    Returns zeros_like val, but can handle arrays and dictionaries.
 
     Parameters
     ----------
-    val : pytree
+    val : PyTree
 
     Returns
     -------
-    zeros_like_val : pytree
+    zeros_like_val : PyTree
         Pytree with the same structure as val with all elements equal to zero.
     """
+
+    # The @partial will force the PyTree to look array-like so we can use array
+    # functions as done below. Using an explicit cast here so LSPs will not complain.
+    val = cast(npt.NDArray, val)
+
     return onp.zeros(val.shape, dtype=val.dtype)
 
 
-def egradRobust(g, j=0):
+def egradRobust(g: Callable[..., Any], j: uint = 0):
     """This function mimics egrad from autograd, but can also handle dictionaries.
 
     Parameters
@@ -151,7 +181,7 @@ def egradRobust(g, j=0):
     if g.__qualname__ == "jit.<locals>.f_jitted":
         g = g.__wrapped__
 
-    def wrapped(*args):
+    def wrapped(*args: Any) -> Any:
         """
         Wrapper for derivative of g with respect to parameter number j.
 
@@ -174,28 +204,28 @@ def wrapped(*args):
     return wrapped
 
 
-def pe(*args: Any, constant_arg_nums: List[int] = ()) -> Any:
+def pe(*args: Any, constant_arg_nums: List[int] = []) -> Any:
     """
     Decorator that returns a function evaluated such that the arg numbers specified in constant_arg_nums
     and all functions that utilizes only those arguments are treated as compile time constants.
 
-    Parameters:
-    -----------
-    *args: Any
+    Parameters
+    ----------
+    *args : Any
         Arguments for the function that pe is applied to.
-    constant_arg_nums: List[int], optional
+    constant_arg_nums : List[int], optional
         The arguments whose values and functions that depend only on these values should be
         treated as cached constants.
 
-    Returns:
-    --------
-    f: Any
+    Returns
+    -------
+    f : Any
         The new function whose constant_arg_num arguments have been removed. The jaxpr of this
         function has the constant_arg_num values and all functions that depend on those values
         cached as constants.
 
-    Usage:
-    ------
+    Usage
+    -----
     @pe(*args, constant_arg_nums=[0])
     def f(x,xi):
         # Function stuff here
@@ -268,7 +298,7 @@ def get_arg(a, unknown):
     return wrapper
 
 
-def pejit(*args: Any, constant_arg_nums: List[int] = (), **kwargs) -> Any:
+def pejit(*args: Any, constant_arg_nums: List[int] = [], **kwargs) -> Any:
     """
     Works like :func:`pe <tfc.utils.TFCUtils.pe>`, but also JITs the returned function. See :func:`pe <tfc.utils.TFCUtils.pe>` for more details.
 
@@ -306,8 +336,15 @@ class TFCDict(OrderedDict):
       - Turning a 1-D array into a dictionary
     """
 
-    def __init__(self, *args):
-        """Initialize TFCDict using the OrderedDict method."""
+    def __init__(self, *args: Any):
+        """
+        Initialize TFCDict using the OrderedDict method.
+
+        Parameters
+        ----------
+        *args : Any
+            Arguments to pass to the OrderedDict
+        """
 
         # Store dictionary and keep a record of the keys. Keys will stay in same
         # order, so that adding and subtracting is repeatable.
@@ -316,8 +353,10 @@ def __init__(self, *args):
         self._nKeys = len(self._keys)
         self.getSlices()
 
-    def getSlices(self):
-        """Function that creates slices for each of the keys in the dictionary."""
+    def getSlices(self) -> None:
+        """
+        Function that creates slices for each of the keys in the dictionary.
+        """
         if all(
             isinstance(value, np.ndarray) or isinstance(value, onp.ndarray)
             for value in self.values()
@@ -338,47 +377,68 @@ def getSlices(self):
                 None,
             ] * self._nKeys
 
-    def update(self, *args):
-        """Overload the update method to update the _keys variable as well.
+    def update(self, *args: Any) -> None:
+        """
+        Overload the update method to update the _keys variable as well.
 
         Parameters
         ----------
-        *args : iterable
-            Same as *args for the update method on ordered dict.
+        *args : Any
+            Same as *args for the update method of ordered dict.
         """
         super().update(*args)
         self._keys = list(self.keys())
         self._nKeys = len(self._keys)
         self.getSlices()
 
-    def toArray(self):
-        """Send dictionary to a flat JAX array."""
-        return np.hstack([k for k in self.values()])
+    def toArray(self) -> np.ndarray:
+        """
+        Send dictionary to a flat JAX array.
 
-    def toDict(self, arr):
-        """Send a flat JAX array to a TFCDict with the same keys.
+        Returns
+        -------
+        np.ndarray
+            This dictionary as a flat JAX array.
+        """
+        return cast(np.ndarray, np.hstack([k for k in self.values()]))
+
+    def toDict(self, arr: np.ndarray) -> "TFCDict":
+        """
+        Send a flat JAX array to a TFCDict with the same keys.
 
         Parameters
         ----------
-        arr : array-like
+        arr : ndarray
             Flat JAX array to convert to TFCDict. Must have the same number of elements as total number of elements in the dictionary.
+
+        Returns
+        -------
+        TFCDict
+            JAX array as a TFCDict
         """
         arr = arr.flatten()
         return TFCDict(zip(self._keys, [arr[self._slices[k]] for k in range(self._nKeys)]))
 
-    def block_until_ready(self):
-        """Mimics block_until_ready for jax arrays. Used to halt the program until the computation that created the
+    def block_until_ready(self) -> "TFCDict":
+        """
+        Mimics block_until_ready for jax arrays. Used to halt the program until the computation that created the
         dictionary is finished.
+
+        Returns
+        -------
+        TFCDict
+            This TFCDict.
         """
         self[self._keys[0]].block_until_ready()
         return self
 
-    def __iadd__(self, o):
-        """Used to overload "+=" for TFCDict so that 2 TFCDict's can be added together.
+    def __iadd__(self, o: TFCDictAddable) -> "TFCDict":
+        """
+        Used to overload "+=" for TFCDict so that 2 TFCDict's can be added together.
 
         Parameters
         ----------
-        o : array-like or Python dictionary or TFCDict
+        o : TFCDictAddable
             Values to add to the current dicitonary.
 
         Returns
@@ -395,16 +455,17 @@ def __iadd__(self, o):
                 self[self._keys[k]] += o[self._slices[k]]
         return self
 
-    def __isub__(self, o):
-        """Used to overload "-=" for TFCDict so that 2 TFCDict's can be subtracted.
+    def __isub__(self, o: TFCDictAddable) -> "TFCDict":
+        """
+        Used to overload "-=" for TFCDict so that 2 TFCDict's can be subtracted.
 
         Parameters
         ----------
-        o : array-like or Python dictionary or TFCDict
+        o : TFCDictAddable
             Values to subtract from the current dicitonary.
 
         Returns
-        ----------
+        -------
         self : TFCDict
             A copy of self after subtracting the values from o.
         """
@@ -417,12 +478,13 @@ def __isub__(self, o):
                 self[self._keys[k]] -= o[self._slices[k]]
         return self
 
-    def __add__(self, o):
-        """Used to overload "+" for TFCDict so that 2 TFCDict's can be added together.
+    def __add__(self, o: TFCDictAddable) -> "TFCDict":
+        """
+        Used to overload "+" for TFCDict so that 2 TFCDict's can be added together.
 
         Parameters
         ----------
-        o : array-like or Python dictionary or TFCDict
+        o : TFCDictAddable
             Values to add to the current dicitonary.
 
         Returns
@@ -440,12 +502,13 @@ def __add__(self, o):
                 out[self._keys[k]] += o[self._slices[k]]
         return out
 
-    def __sub__(self, o):
-        """Used to overload "-" for TFCDict so that 2 TFCDict's can be subtracted.
+    def __sub__(self, o: TFCDictAddable) -> "TFCDict":
+        """
+        Used to overload "-" for TFCDict so that 2 TFCDict's can be subtracted.
 
         Parameters
         ----------
-        o : array-like or Python dictionary or TFCDict
+        o : TFCDictAddable
             Values to subtract from the current dicitonary.
 
         Returns
@@ -475,8 +538,15 @@ def __sub__(self, o):
 class TFCDictRobust(OrderedDict):
     """This class is like the :class:`TFCDict <tfc.utils.TFCUtils.TFCDict>` class, but it handles non-flat arrays."""
 
-    def __init__(self, *args):
-        """Initialize TFCDictRobust using the OrderedDict method."""
+    def __init__(self, *args: Any):
+        """
+        Initialize TFCDictRobust using the OrderedDict method.
+
+        Parameters
+        ----------
+        *args : Any
+            Arguments to pass to the OrderedDict
+        """
 
         # Store dictionary and keep a record of the keys. Keys will stay in same
         # order, so that adding and subtracting is repeatable.
@@ -485,8 +555,10 @@ def __init__(self, *args):
         self._nKeys = len(self._keys)
         self.getSlices()
 
-    def getSlices(self):
-        """Function that creates slices for each of the keys in the dictionary."""
+    def getSlices(self) -> None:
+        """
+        Function that creates slices for each of the keys in the dictionary.
+        """
         if all(isinstance(value, np.ndarray) for value in self.values()):
             arrLen = 0
             self._slices = [
@@ -504,12 +576,13 @@ def getSlices(self):
                 None,
             ] * self._nKeys
 
-    def update(self, *args):
-        """Overload the update method to update the _keys variable as well.
+    def update(self, *args: Any) -> None:
+        """
+        Overload the update method to update the _keys variable as well.
 
         Parameters
         ----------
-        *args : iterable
+        *args : Any
             Same as *args for the update method on ordered dict.
         """
         super().update(*args)
@@ -517,17 +590,30 @@ def update(self, *args):
         self._nKeys = len(self._keys)
         self.getSlices()
 
-    def toArray(self):
-        """Send dictionary to a flat JAX array."""
-        return np.hstack([k.flatten() for k in self.values()])
+    def toArray(self) -> np.ndarray:
+        """
+        Send dictionary to a flat JAX array.
 
-    def toDict(self, arr):
-        """Send a flat JAX array to a TFCDictRobust with the same keys.
+        Returns
+        -------
+        np.ndarray
+            This dictionary as a flat JAX array.
+        """
+        return cast(np.ndarray, np.hstack([k.flatten() for k in self.values()]))
+
+    def toDict(self, arr: np.ndarray) -> "TFCDictRobust":
+        """
+        Send a flat JAX array to a TFCDictRobust with the same keys.
 
         Parameters
         ----------
-        arr : array-like
+        arr : np.ndarray
             Flat JAX array to convert to TFCDictRobust. Must have the same number of elements as total number of elements in the dictionary.
+
+        Returns
+        -------
+        TFCDictRobust
+            JAX array as a TFCDictRobust
         """
         arr = arr.flatten()
         return TFCDictRobust(
@@ -540,19 +626,26 @@ def toDict(self, arr):
             )
         )
 
-    def block_until_ready(self):
-        """Mimics block_until_ready for jax arrays. Used to halt the program until the computation that created the
+    def block_until_ready(self) -> "TFCDictRobust":
+        """
+        Mimics block_until_ready for jax arrays. Used to halt the program until the computation that created the
         dictionary is finished.
+
+        Returns
+        -------
+        TFCDictRobust
+            This TFCDictRobust
         """
         self[self._keys[0]].block_until_ready()
         return self
 
-    def __iadd__(self, o):
-        """Used to overload "+=" for TFCDictRobust so that 2 TFCDictRobust's can be added together.
+    def __iadd__(self, o: TFCDictRobustAddable) -> "TFCDictRobust":
+        """
+        Used to overload "+=" for TFCDictRobust so that 2 TFCDictRobust's can be added together.
 
         Parameters
         ----------
-        o : array-like or Python dictionary or TFCDictRobust
+        o : TFCDictRobustAddable
             Values to add to the current dicitonary.
 
         Returns
@@ -569,12 +662,13 @@ def __iadd__(self, o):
                 self[self._keys[k]] += o[self._slices[k]].reshape(self[self._keys[k]].shape)
         return self
 
-    def __isub__(self, o):
-        """Used to overload "-=" for TFCDictRobust so that 2 TFCDictRobust's can be subtracted.
+    def __isub__(self, o: TFCDictRobustAddable) -> "TFCDictRobust":
+        """
+        Used to overload "-=" for TFCDictRobust so that 2 TFCDictRobust's can be subtracted.
 
         Parameters
         ----------
-        o : array-like or Python dictionary or TFCDictRobust
+        o : TFCDictRobustAddable
             Values to subtract from the current dicitonary.
 
         Returns
@@ -591,12 +685,13 @@ def __isub__(self, o):
                 self[self._keys[k]] -= o[self._slices[k]].reshape(self[self._keys[k]].shape)
         return self
 
-    def __add__(self, o):
-        """Used to overload "+" for TFCDictRobust so that 2 TFCDictRobust's can be added together.
+    def __add__(self, o: TFCDictRobustAddable) -> "TFCDictRobust":
+        """
+        Used to overload "+" for TFCDictRobust so that 2 TFCDictRobust's can be added together.
 
         Parameters
         ----------
-        o : array-like or Python dictionary or TFCDictRobust
+        o : TFCDictRobustAddable
             Values to add to the current dicitonary.
 
         Returns
@@ -614,16 +709,17 @@ def __add__(self, o):
                 out[self._keys[k]] += o[self._slices[k]].reshape(self[self._keys[k]].shape)
         return out
 
-    def __sub__(self, o):
-        """Used to overload "-" for TFCDictRobust so that 2 TFCDictRobust's can be subtracted.
+    def __sub__(self, o: TFCDictRobustAddable) -> "TFCDictRobust":
+        """
+        Used to overload "-" for TFCDictRobust so that 2 TFCDictRobust's can be subtracted.
 
         Parameters
         ----------
-        o : array-like or Python dictionary or TFCDictRobust
+        o : TFCDictRobustAddable
             Values to subtract from the current dicitonary.
 
         Returns
-        ----------
+        -------
         self : TFCDictRobust
             A TFCDictRobust with values = self - o.
         """
@@ -647,16 +743,16 @@ def __sub__(self, o):
 
 
 def LS(
-    zXi,
-    res,
-    *args,
+    zXi: PyTree,
+    res: Callable,
+    *args: Any,
     constant_arg_nums: List[int] = [],
-    J=None,
-    method="pinv",
-    timer=False,
-    timerType="process_time",
-    holomorphic=False,
-):
+    J: Optional[Callable[..., np.ndarray]] = None,
+    method: Literal["pinv", "lstsq"] = "pinv",
+    timer: bool = False,
+    timerType: str = "process_time",
+    holomorphic: bool = False,
+) -> Union[PyTree, Tuple[PyTree, float]]:
     """
     JITed least squares.
     This function takes in an initial guess of zeros, zXi, and a residual function, res, and
@@ -665,22 +761,23 @@ def LS(
 
     Parameters
     ----------
-    zXi : pytree or array-like
+    zXi : PyTree
         Unknown parameters to be found using least-squares.
 
-    res : function
-        Residual function (also known as the loss function) with signature res(xi,*args).
+    res : Callable
+        Residual function (also known as the loss function) with signature res(xi: PyTree, *args:Any, **kwargs:Any).
+        Note, the first argument does not need to be named xi, this is just illustrative.
 
-    *args : iterable
-        Any additional arguments taken by res other than xi.
+    *args : Any
+        Any additional arguments taken by res other than the first PyTree argument.
 
     constant_arg_nums: List[int], optional
         These arguments will be removed from the residual function and treated as constant. See :func:`pejit <tfc.utils.TFCUtils.pejit>` for more details.
 
-    J : function, optional
-         User specified Jacobian. If None, then the Jacobian of res with respect to xi will be calculated via automatic differentiation. (Default value = None)
+    J : Optional[Callable[...,np.ndarray]]
+         User specified Jacobian function. If None, then the Jacobian of res with respect to xi will be calculated via automatic differentiation. (Default value = None)
 
-    method : {"pinv","lstsq"}, optional
+    method : Literal["pinv","lstsq"], optional
          Method for least-squares inversion. (Default value = "pinv")
          * pinv - Use np.linalg.pinv
          * lstsq - Use np.linalg.lstsq
@@ -742,7 +839,7 @@ def J(xi, *args):
         # Make arguments constant if desired
         ls = pe(zXi, *args, constant_arg_nums=constant_arg_nums)(ls)
 
-        args = list(args)
+        args: List[Any] = list(args)
         constant_arg_nums.sort()
         constant_arg_nums.reverse()
         for k in constant_arg_nums:
@@ -754,12 +851,12 @@ def J(xi, *args):
     if timer:
         import time
 
-        timer = getattr(time, timerType)
+        timer_f: Callable[[], float] = getattr(time, timerType)
         ls(zXi, *args).block_until_ready()
 
-        start = timer()
+        start = timer_f()
         xi = ls(zXi, *args).block_until_ready()
-        stop = timer()
+        stop = timer_f()
         zXi += xi
 
         return zXi, stop - start
@@ -777,17 +874,52 @@ class LsClass:
 
     def __init__(
         self,
-        zXi,
-        res,
-        *args,
-        J=None,
+        zXi: PyTree,
+        res: Callable,
+        *args: Any,
         constant_arg_nums: List[int] = [],
-        method="pinv",
-        timer=False,
-        timerType="process_time",
-        holomorphic=False,
-    ):
-        """Initialization function. Creates the JIT-ed least-squares function."""
+        J: Optional[Callable[..., np.ndarray]] = None,
+        method: Literal["pinv", "lstsq"] = "pinv",
+        timer: bool = False,
+        timerType: str = "process_time",
+        holomorphic: bool = False,
+    ) -> None:
+        """
+        Initialization function. Creates the JIT-ed least-squares function.
+
+        Parameters
+        ----------
+        zXi : PyTree
+            Unknown parameters to be found using least-squares.
+
+        res : Callable
+            Residual function (also known as the loss function) with signature res(xi: PyTree, *args:Any, **kwargs:Any).
+            Note, the first argument does not need to be named xi, this is just illustrative.
+
+        *args : Any
+            Any additional arguments taken by res other than the first PyTree argument.
+
+        J : Optional[Callable[...,np.ndarray]]
+             User specified Jacobian function. If None, then the Jacobian of res with respect to xi will be calculated via automatic differentiation. (Default value = None)
+
+        constant_arg_nums: List[int], optional
+            These arguments will be removed from the residual function and treated as constant. See :func:`pejit <tfc.utils.TFCUtils.pejit>` for more details.
+
+        method : Literal["pinv","lstsq"], optional
+             Method for least-squares inversion. (Default value = "pinv")
+             * pinv - Use np.linalg.pinv
+             * lstsq - Use np.linalg.lstsq
+
+        timer : bool, optional
+             Boolean that chooses whether to time the code or not. (Default value = False). Note that setting to true adds a slight increase in runtime.
+             As one iteration of the non-linear least squares is run first to avoid timining the JAX trace.
+
+        timerType : str, optional
+             Any timer from the time module. (Default value = "process_time")
+
+        holomorphic : bool, optional
+             Indicates whether residual function is promised to be holomorphic. (Default value = False)
+        """
 
         self.timerType = timerType
         self.timer = timer
@@ -834,7 +966,7 @@ def J(xi, *args):
             # Make arguments constant if desired
             ls = pe(zXi, *args, constant_arg_nums=constant_arg_nums)(ls)
 
-            args = list(args)
+            args: List[Any] = list(args)
             constant_arg_nums.sort()
             constant_arg_nums.reverse()
             for k in constant_arg_nums:
@@ -844,23 +976,24 @@ def J(xi, *args):
 
         self._compiled = False
 
-    def run(self, zXi, *args):
-        """Runs the JIT-ed least-squares function and times it if desired.
+    def run(self, zXi: PyTree, *args: Any) -> Union[PyTree, Tuple[PyTree, float]]:
+        """
+        Runs the JIT-ed least-squares function and times it if desired.
 
         Parameters
         ----------
-        zXi : pytree or array-like
+        zXi : PyTree
             Unknown parameters to be found using least-squares.
 
-        *args : iterable
+        *args : Any
             Any additional arguments taken by res other than xi.
 
         Returns
         -------
-        xi : pytree or array-like
+        xi : PyTree
              Unknowns that minimize res as found via least-squares. Type will be the same as zXi specified in the input.
 
-        time : float
+        time : float, optional
              Computation time as calculated by timerType specified. This output is only returned if timer = True.
 
         """
@@ -889,33 +1022,33 @@ def run(self, zXi, *args):
             return zXi
 
 
-def nlls_printout(arg, transforms, *, end="\n", **kwargs):
+def nlls_printout(arg: Any, transforms: Any, *, end: str = "\n", **kwargs: Any) -> None:
     print("Iteration: {0}\tmax(abs(res)): {1}".format(*arg), end=end)
     return
 
 
-def nlls_id_print(it, x, end="\n"):
+def nlls_id_print(it: int, x, end: str = "\n"):
     printer = partial(nlls_printout, end=end)
     return id_tap(printer, (it, x))
 
 
 def NLLS(
-    xiInit,
-    res,
-    *args,
+    xiInit: PyTree,
+    res: Callable,
+    *args: Any,
     constant_arg_nums: List[int] = [],
-    J=None,
-    cond=None,
-    body=None,
-    tol=1e-13,
-    maxIter=50,
-    method="pinv",
-    timer=False,
-    printOut=False,
-    printOutEnd="\n",
-    timerType="process_time",
-    holomorphic=False,
-):
+    J: Optional[Callable[..., np.ndarray]] = None,
+    cond: Optional[Callable[[PyTree], bool]] = None,
+    body: Optional[Callable[[PyTree], PyTree]] = None,
+    tol: float = 1e-13,
+    maxIter: uint = 50,
+    method: Literal["pinv", "lstsq"] = "pinv",
+    timer: bool = False,
+    printOut: bool = False,
+    printOutEnd: str = "\n",
+    timerType: str = "process_time",
+    holomorphic: bool = False,
+) -> Union[Tuple[PyTree, int], Tuple[PyTree, int, float]]:
     """
     JIT-ed non-linear least squares.
     This function takes in an initial guess, xiInit (initial values of xi), and a residual function, res, and
@@ -942,11 +1075,11 @@ def NLLS(
     J : function
          User specified Jacobian. If None, then the Jacobian of res with respect to xi will be calculated via automatic differentiation. (Default value = None)
 
-    cond : function, optional
+    cond : Optional[Callable[[PyTree],bool]]
          User specified condition function. If None, then the default cond function is used which checks the three termination criteria
          provided in the class description. (Default value = None)
 
-    body : function, optional
+    body : Optional[Callable[[PyTree],PyTree]]
          User specified while-loop body function. If None, then use the default body function which updates xi using a NLLS interation and the method provided.
          (Default value = None)
 
@@ -956,7 +1089,7 @@ def NLLS(
     maxIter : int, optional
          Maximum number of iterations. (Default value = 50)
 
-    method : {"pinv","lstsq"}, optional
+    method : Literal["pinv","lstsq"], optional
          Method for least-squares inversion. (Default value = "pinv")
          * pinv - Use np.linalg.pinv
          * lstsq - Use np.linalg.lstsq
@@ -979,11 +1112,11 @@ def NLLS(
 
     Returns
     -------
-    xi : pytree or array-like
+    xi : PyTree
          Unknowns that minimize res as found via least-squares. Type will be the same as zXi specified in the input.
 
     it : int
-         Number of NLLS iterations performed..
+         Number of NLLS iterations performed.
 
     time : float
          Computation time as calculated by timerType specified. This output is only returned if timer = True.
@@ -1044,7 +1177,7 @@ def J(xi, *args):
         LS = pe(xiInit, *args, constant_arg_nums=constant_arg_nums)(LS)
         res = pe(xiInit, *args, constant_arg_nums=constant_arg_nums)(res)
 
-        args = list(args)
+        args: List[Any] = list(args)
         constant_arg_nums.sort()
         constant_arg_nums.reverse()
         for k in constant_arg_nums:
@@ -1073,23 +1206,23 @@ def body(val):
     nlls = jit(lambda val: lax.while_loop(cond, body, val))
 
     if dictFlag:
-        dxi = np.ones_like(xiInit.toArray())
+        dxi = np.ones_like(cast(Union[TFCDict, TFCDictRobust], xiInit).toArray())
     else:
         dxi = np.ones_like(xiInit)
 
     if timer:
         import time
 
-        timer = getattr(time, timerType)
+        timer_f: Callable[[], float] = getattr(time, timerType)
         val = {"xi": xiInit, "dxi": dxi, "it": maxIter - 1, "args": args}
         nlls(val)["dxi"].block_until_ready()
 
         val = {"xi": xiInit, "dxi": dxi, "it": 0, "args": args}
 
-        start = timer()
+        start = timer_f()
         val = nlls(val)
         val["dxi"].block_until_ready()
-        stop = timer()
+        stop = timer_f()
 
         return val["xi"], val["it"], stop - start
     else:
@@ -1106,23 +1239,77 @@ class NllsClass:
 
     def __init__(
         self,
-        xiInit,
-        res,
-        *args,
+        xiInit: PyTree,
+        res: Callable,
+        *args: Any,
         constant_arg_nums: List[int] = [],
-        J=None,
-        cond=None,
-        body=None,
-        tol=1e-13,
-        maxIter=50,
-        method="pinv",
-        timer=False,
-        printOut=False,
-        printOutEnd="\n",
-        timerType="process_time",
-        holomorphic=False,
-    ):
-        """Initialization function. Creates the JIT-ed nonlinear least-squares function."""
+        J: Optional[Callable[..., np.ndarray]] = None,
+        cond: Optional[Callable[[PyTree], bool]] = None,
+        body: Optional[Callable[[PyTree], PyTree]] = None,
+        tol: float = 1e-13,
+        maxIter: uint = 50,
+        method: Literal["pinv", "lstsq"] = "pinv",
+        timer: bool = False,
+        printOut: bool = False,
+        printOutEnd: str = "\n",
+        timerType: str = "process_time",
+        holomorphic: bool = False,
+    ) -> None:
+        """
+        Initialization function. Creates the JIT-ed nonlinear least-squares function.
+
+        Parameters
+        ----------
+        xiInit : pytree or array-like
+            Initial guess for the unkown parameters.
+
+        res : function
+            Residual function (also known as the loss function) with signature res(xi,*args).
+
+        *args : iterable
+            Any additional arguments taken by res other than xi.
+
+        constant_arg_nums: List[int], optional
+            These arguments will be removed from the residual function and treated as constant. See :func:`pejit <tfc.utils.TFCUtils.pejit>` for more details.
+
+        J : function
+             User specified Jacobian. If None, then the Jacobian of res with respect to xi will be calculated via automatic differentiation. (Default value = None)
+
+        cond : Optional[Callable[[PyTree],bool]]
+             User specified condition function. If None, then the default cond function is used which checks the three termination criteria
+             provided in the class description. (Default value = None)
+
+        body : Optional[Callable[[PyTree],PyTree]]
+             User specified while-loop body function. If None, then use the default body function which updates xi using a NLLS interation and the method provided.
+             (Default value = None)
+
+        tol : float
+             Tolerance used in the default termination criteria: see class description for more details. (Default value = 1e-13)
+
+        maxIter : int, optional
+             Maximum number of iterations. (Default value = 50)
+
+        method : Literal["pinv","lstsq"], optional
+             Method for least-squares inversion. (Default value = "pinv")
+             * pinv - Use np.linalg.pinv
+             * lstsq - Use np.linalg.lstsq
+
+        timer : bool, optional
+             Boolean that chooses whether to time the code or not. (Default value = False). Note that setting to true adds a slight increase in runtime.
+             As one iteration of the non-linear least squares is run first to avoid timining the JAX trace.
+
+        printOut : bool, optional
+             Controls whether the NLLS prints out information each interaton or not. The printout consists of the iteration and max(abs(res)) at each iteration. (Default value = False)
+
+        printOutEnd : str, optional
+             Value of keyword argument end passed to the print statement used in printOut. (Default value = "\\\\n")
+
+        timerType : str, optional
+             Any timer from the time module. (Default value = "process_time")
+
+        holomorphic : bool, optional
+             Indicates whether residual function is promised to be holomorphic. (Default value = False)
+        """
 
         self.timerType = timerType
         self.timer = timer
@@ -1186,7 +1373,7 @@ def J(xi, *args):
             LS = pe(xiInit, *args, constant_arg_nums=constant_arg_nums)(LS)
             res = pe(xiInit, *args, constant_arg_nums=constant_arg_nums)(res)
 
-            args = list(args)
+            args: List[Any] = list(args)
             constant_arg_nums.sort()
             constant_arg_nums.reverse()
             for k in constant_arg_nums:
@@ -1215,20 +1402,22 @@ def body(val):
         self._nlls = jit(lambda val: lax.while_loop(cond, body, val))
         self._compiled = False
 
-    def run(self, xiInit, *args):
+    def run(
+        self, xiInit: PyTree, *args: Any
+    ) -> Union[Tuple[PyTree, int], Tuple[PyTree, int, float]]:
         """Runs the JIT-ed nonlinear least-squares function and times it if desired.
 
         Parameters
         ----------
-        xiInit : pytree or array-like
+        xiInit : PyTree
             Initial guess for the unkown parameters.
 
-        *args : iterable
+        *args : Any
             Any additional arguments taken by res other than xi.
 
         Returns
         -------
-        xi : pytree or array-like
+        xi : PyTree
              Unknowns that minimize res as found via least-squares. Type will be the same as zXi specified in the input.
 
         it : int
@@ -1240,14 +1429,14 @@ def run(self, xiInit, *args):
         """
 
         if self._dictFlag:
-            dxi = np.ones_like(xiInit.toArray())
+            dxi = np.ones_like(cast(Union[TFCDict, TFCDictRobust], xiInit).toArray())
         else:
             dxi = np.ones_like(xiInit)
 
         if self.timer:
             import time
 
-            timer = getattr(time, self.timerType)
+            timer_f: Callable[[], float] = getattr(time, self.timerType)
 
             if not self._compiled:
                 val = {"xi": xiInit, "dxi": dxi, "it": self._maxIter - 1, "args": args}
@@ -1256,10 +1445,10 @@ def run(self, xiInit, *args):
 
             val = {"xi": xiInit, "dxi": dxi, "it": 0, "args": args}
 
-            start = timer()
+            start = timer_f()
             val = self._nlls(val)
             val["dxi"].block_until_ready()
-            stop = timer()
+            stop = timer_f()
 
             return val["xi"], val["it"], stop - start
 
@@ -1272,32 +1461,28 @@ def run(self, xiInit, *args):
             return val["xi"], val["it"]
 
 
+class ComponentConstraintDict(TypedDict):
+    name: str
+    node0: str
+    node1: str
+
+
 class ComponentConstraintGraph:
     """
     Creates a graph of all valid ways in which component constraints can be embedded.
-
-    Parameters
-    ----------
-    N : list
-        A list of strings that specify the node names. These node names typically coincide with
-        the names of the dependent variables.
-    E : dict
-        A dictionary with the following fields:
-        * name - Name of the component constraint.
-        * node0 - The name of one of the nodes that makes up the component constraint.  Must correspond with an element of the list given in N.
-        * node1 - The name of one of the nodes that makes up the component constraint.  Must correspond with an element of the list given in N.
     """
 
-    def __init__(self, N, E):
-        """Class constructor.
+    def __init__(self, N: List[str], E: List[ComponentConstraintDict]) -> None:
+        """
+        Class constructor.
 
         Parameters
         ----------
-        N : list
+        N : List[str]
             A list of strings that specify the node names. These node names typically coincide with
             the names of the dependent variables.
-        E : dict
-            A dictionary with the following fields:
+        E : List[ComponentConstraintDict]
+            The ComponentConstraintDict is a dictionary with the following fields:
             * name - Name of the component constraint.
             * node0 - The name of one of the nodes that makes up the component constraint.  Must correspond with an element of the list given in N.
             * node1 - The name of one of the nodes that makes up the component constraint.  Must correspond with an element of the list given in N.
@@ -1324,12 +1509,9 @@ def __init__(self, N, E):
         # Find all targets that are valid trees
         self.goodTargets = []
         for j in range(len(self.targets)):
-            flag = True
             adj = onp.zeros((self.nNodes, self.nNodes))
             for k in range(self.nNodes):
                 kNode = N[k]
-                sources = []
-                targets = []
                 for g in range(self.nEdges):
                     if E[g]["node0"] == kNode:
                         if self.targets[j][g]:
@@ -1344,14 +1526,14 @@ def __init__(self, N, E):
         self.N = N
         self.E = E
 
-    def SaveGraphs(self, outputDir, allGraphs=False, savePDFs=False):
+    def SaveGraphs(self, outputDir: Path, allGraphs: bool = False, savePDFs: bool = False) -> None:
         """
         Saves the graphs.
         The graphs are saved in a clickable HTML structure. They can also be saved as PDFs.
 
         Parameters
         ----------
-        outputDir : str
+        outputDir : Path
             Output directory to save in.
 
         allGraphs : bool, optional
@@ -1447,62 +1629,63 @@ def SaveGraphs(self, outputDir, allGraphs=False, savePDFs=False):
             treeHtml.WriteFile()
 
 
-def ScaledQrLs(A, B):
+def ScaledQrLs(A: np.ndarray, B: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
     """This function performs least-squares using a scaled QR method.
 
     Parameters
     ----------
-    A : array-like
+    A : np.ndarray
         A matrix in A*x = B.
 
-    B : array-like
+    B : np.ndarray
         B matrix in A*x = B.
 
     Returns
     -------
-    x : array-like
+    x : np.ndarray
         Solution to A*x = B solved using a scaled QR method.
 
-    cn : array-like
+    cn : np.ndarray
         Condition number.
     """
     S = 1.0 / np.sqrt(np.sum(A * A, 0))
     S = np.reshape(S, (A.shape[1],))
     q, r = np.linalg.qr(A.dot(np.diag(S)))
-    x = S * np.linalg.multi_dot([_MatPinv(r), q.T, B])
-    cn = np.linalg.cond(r)
+    x: np.ndarray = S * np.linalg.multi_dot([_MatPinv(r), q.T, B])
+    cn: np.ndarray = cast(np.ndarray, np.linalg.cond(r))
     return x, cn
 
 
-def _MatPinv(A):
+def _MatPinv(A: np.ndarray) -> np.ndarray:
     """This function is used to better replicate MATLAB's pseudo-inverse.
 
     Parameters
     ----------
-    A : array-like
+    A : np.ndarray
         Matrix to be inverted.
 
     Returns
     -------
-    Ainv : array-like
+    Ainv : np.ndarray
         Inverse of A.
     """
     rcond = onp.max(A.shape) * onp.spacing(np.linalg.norm(A, ord=2))
     return np.linalg.pinv(A, rcond=rcond)
 
 
-def step(x):
-    """This is the unit step function, but the deriative is defined and equal to 0 at every point.
+def step(x: np.ndarray) -> np.ndarray:
+    """
+    This is the unit step function, but the deriative is defined and equal to 0 at every point.
 
     Parameters
     ----------
-    x : array-like
+    x : np.ndarray
         Array to apply step to.
 
 
     Returns
     -------
-    step_x : array-like
+    step_x : np.ndarray
         step(x)
     """
     return np.heaviside(x, 0)
diff --git a/src/tfc/utils/types.py b/src/tfc/utils/types.py
index 1271eb0..8b4970b 100644
--- a/src/tfc/utils/types.py
+++ b/src/tfc/utils/types.py
@@ -1,21 +1,65 @@
 import sys
-from typing import Union
+from typing import Union, Any
+import numpy as np
+import numpy.typing as npt
+import jax.numpy as jnp
+from jaxtyping import Array
 
 if sys.version_info >= (3, 8):
-    from typing import Literal
-
-    Literal = Literal
+    from typing import Literal, Protocol, TypedDict
 else:
-    from typing_extensions import Literal
-
-    Literal = Literal
+    from typing_extensions import Literal, Protocol, TypedDict
 
 if sys.version_info >= (3, 9):
     from typing import Annotated
+
+    List = list
+    Tuple = tuple
+    Dict = dict
+    Tuple = tuple
 else:
     from typing_extensions import Annotated
+    from typing import List, Tuple, Dict, Tuple
+
+from annotated_types import Gt, Ge, Lt, Le
 
-from annotated_types import Gt
+# Path
+# Path = Union[str, os.PathLike]
+Path = str
 
-uint = Annotated[int, Gt(0)]
+# Integer > 0
+pint = Annotated[int, Gt(0)]
+
+# Integer >= 0
+uint = Annotated[int, Ge(0)]
+
+# General number type
 Number = Union[int, float, complex]
+
+from numpy._typing._array_like import _ArrayLikeStr_co, _ArrayLikeInt_co
+
+# Array-like of strings
+StrArrayLike = _ArrayLikeStr_co
+
+# Array-like of integers
+IntArrayLike = _ArrayLikeInt_co
+
+# List or array like
+NumberListOrArray = Union[Tuple[Number, ...], List[Number], npt.NDArray[Any], Array]
+
+# List or array of integers
+IntListOrArray = Union[
+    Tuple[int, ...],
+    List[int],
+    npt.NDArray[np.int32],
+    npt.NDArray[np.int64],
+    npt.NDArray[np.int16],
+    npt.NDArray[np.int8],
+]
+
+# JAX array or numpy array
+JaxOrNumpyArray = Union[npt.NDArray, Array]
+
+# Tuple or list of array
+TupleOrListOfArray = Union[Tuple[JaxOrNumpyArray, ...], List[JaxOrNumpyArray]]
+TupleOrListOfNumpyArray = Union[Tuple[npt.NDArray, ...], List[npt.NDArray]]


Test patch:
diff --git a/tests/test_all b/tests/test_all
index c9d5456..2224c3c 100755
--- a/tests/test_all
+++ b/tests/test_all
@@ -23,7 +23,10 @@ while [[ $# -gt 0 ]]; do
 done
 
 # Collect all of the test cases
-ALL_TESTS=`find ./temp -name "*.py" | xargs grep -Li tensorflow`
+# SKIP_TESTS are test cases that don't have tensorflow in their imports, but
+# use files that do.
+ALL_TESTS=`find ../examples -name "*.py" | xargs grep -Li tensorflow`
+SKIP_TESTS=("../examples/Carl_Leake_Dissertation/Chapter_3/Example_3_4.py"
 CD=`pwd`
 
 # Function to run a test
@@ -44,9 +47,13 @@ num_jobs="\j"
 
 # Run all of the jobs
 for TEST in ${ALL_TESTS[@]}; do
-	while (( ${num_jobs@P} >= num_procs )); do
-    	wait -n
-    done
-	run_test $TEST &
+	if [[ " ${SKIP_TESTS[*]} " =~ " ${TEST} " ]]; then
+		echo Skipping $TEST
+	else
+		while (( ${num_jobs@P} >= num_procs )); do
+			wait -n
+		done
+		run_test $TEST &
+	fi
 done
 wait