Repository: cair/tmu (language: python)

Issue:
Update

get_encoded_data does not refresh with new data
Had to remove the compute_hash part to make it work:

def get_encoded_data(self, data: np.ndarray, encoder_func) -> np.ndarray:
        """Get encoded data for an array, using cache if available."""

        self.encoded_data = encoder_func(data)

        return self.encoded_data
If all classes are '0' in Y, fit goes into infinite loop
Running vanilla classifier on a dataset where all Y's are 0. Fit never returns.


Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index e0c00141..323d36e2 100644
--- a/README.md
+++ b/README.md
@@ -27,6 +27,9 @@ TMU is a comprehensive repository that encompasses several Tsetlin Machine imple
 
 - Wrappers for C and CUDA-based clause evaluation and updates to enable high-performance computation.
 
+## Guides and Tutorials
+- [Setting up efficient Development Environment](docs/tutorials/devcontainers/devcontainers.md)
+
 ## 📦 Installation
 
 #### **Prerequisites for Windows**
@@ -40,7 +43,11 @@ Ubuntu: `sudo apt install libffi-dev`
 #### **Installing TMU**
 To get started with TMU, run the following command:
 ```bash
+# Installing Stable Branch
 pip install git+https://github.com/cair/tmu.git
+
+# Installing Development Branch
+pip install git+https://github.com/cair/tmu.git@dev
 ```
 
 ## 🛠 Development
@@ -49,22 +56,26 @@ If you're looking to contribute or experiment with the codebase, follow these st
 
 1. **Clone the Repository**:
    ```bash
-   git clone git@github.com:cair/tmu.git
+   git clone -b dev git@github.com:cair/tmu.git && cd tmu
    ```
 
 2. **Set Up Development Environment**:
    Navigate to the project directory and compile the C library:
    ```bash
-   cd tmu && pip install develop .
+   # Install TMU
+    pip install .
+   
+   # (Alternative): Install TMU in Development Mode
+    pip install -e .
+   
+   # Install TMU-Composite
+    pip install .[composite]
+   
+   # Install TMU-Composite in Development Mode
+    pip install -e .[composite]
    ```
 
 3. **Starting a New Project**:
-   For your projects, simply create a new folder within 'examples' and initiate your development.
-
-#### Modifying the C Codebase
-If you make changes to the C codebase, ensure you recompile the code using:
-```bash
-pip install develop .
-```
+   For your projects, simply create a new **branch** and then within the 'examples' folder, create a new project and initiate your development.
 
 ---
diff --git a/docs/tutorials/devcontainers/devcontainers.md b/docs/tutorials/devcontainers/devcontainers.md
new file mode 100644
index 00000000..830b7788
--- /dev/null
+++ b/docs/tutorials/devcontainers/devcontainers.md
@@ -0,0 +1,190 @@
+# Running TMU development in development containers
+
+The recommended way of doing development using TMU is through devcontainers. This is because you get a unified development environment that is independent of your underlying system, making it much more trivial to debug when issues appear, and also is much faster to get up and running.
+This guide shows how to run TMU with VSCODE both in remote development with SSH with devcontainers, and finally with Remote SSH combined with Devcontainers
+
+# 0. Prerequisites
+
+- Docker installed on your local machine. **Optional - When not using remote SSH**
+- Visual Studio Code installed.
+- Remote - Containers extension installed in VSCode.
+- Ensure Git is installed on your system. This guide utilizes Git Bash for command execution across all operating systems, including Windows.
+
+# 1. SSH Configuration
+
+### Step 1: Generate SSH Keys
+
+Regardless of your operating system, the first step is to generate an SSH key pair if you haven't done so already or wish to create a new pair for this connection. Open Git Bash and type:
+
+```bash
+ssh-keygen -t rsa -b 4096 -C "your_email@example.com"
+
+```
+
+Replace `your_email@example.com` with your actual email address for identification purposes. When prompted to "Enter a file in which to save the key," press Enter to accept the default location. You will then be asked to enter a passphrase; you can choose to enter one for added security or press Enter to proceed without a passphrase.
+
+### Step 2: Copy the SSH Key to Your Remote Machine Using `ssh-copy-id`
+
+The `ssh-copy-id` script offers a convenient way to install your public key in a remote machine's `~/.ssh/authorized_keys` file, enabling password-less SSH access. This tool is available on Linux, macOS, and through Git Bash on Windows. Here's how to use it:
+
+### For Linux, macOS, and Windows (Using Git Bash)
+
+1. Open your terminal or Git Bash (for Windows users).
+2. Execute the following command, replacing `<username>` with your user on the remote machine and `cair-gpuXX.uia.no` with the remote machine's hostname or IP address:
+
+```bash
+ssh-copy-id <username>@cair-gpuXX.uia.no
+
+```
+
+1. You will be prompted to enter the remote machine's password. After successfully authenticating, your SSH key will be added to the remote machine's `~/.ssh/authorized_keys` file.
+
+### Step 3: Verify Your SSH Setup
+
+To test your new SSH key setup, attempt to SSH into your remote machine:
+
+```bash
+ssh <username>@cair-gpuXX.uia.no
+
+```
+
+If the setup is correct, you should gain access without being prompted for the remote user's password.
+
+### For Multiple Remote Machines
+
+If you're planning to establish password-less SSH connections with multiple remote machines, simply repeat the above process for each, ensuring your public SSH key is added to each machine's `~/.ssh/authorized_keys` file.
+
+# 2. Development using a Devcontainer
+
+# Step 1: Create the Docker-Compose File
+
+Create a **`docker-compose.yml`** file within your **`.devcontainer`** directory. This file will define your service, including the use of GPUs.
+
+### CUDA Enabled Configuration
+
+**docker-compose.yml** (**`./.devcontainer/docker-compose.yml`**):
+
+```yaml
+version: '3.8'
+services:
+  tmu-development:
+    build:
+      context: .
+      dockerfile: Dockerfile
+    volumes:
+      - ..:/app
+    deploy:
+      resources:
+        reservations:
+          devices:
+          - driver: nvidia
+            capabilities: [gpu]
+            count: 1 # Assign number of GPUs or use 'all' to assign all available GPUs
+```
+
+### CPU-only Enabled Configuration
+
+**docker-compose.yml** (**`./.devcontainer/docker-compose.yml`**):
+
+```yaml
+version: '3.8'
+services:
+  tmu-development:
+    build:
+      context: .
+      dockerfile: Dockerfile
+    volumes:
+      - ..:/app
+
+```
+
+## Step 2: Create a Dockerfile
+
+### CUDA Enabled Configuration
+
+**Dockerfile** (**`./.devcontainer/Dockerfile`**):
+
+```docker
+FROM nvidia/cuda:12.3.2-cudnn9-devel-ubuntu22.04
+
+# Install Python and other dependencies
+RUN apt-get update && apt-get install -y python3-pip
+
+# Install TMU?, other relevant stuff
+
+WORKDIR /app
+COPY . /app
+
+# You should have a requirements.txt to define your dependencies
+RUN pip3 install --no-cache-dir -r requirements.txt
+
+CMD [ "tail", "-f", "/dev/null" ]
+```
+
+### CPU-only Enabled Configuration
+
+**Dockerfile** (**`./.devcontainer/Dockerfile`**):
+
+```docker
+FROM ubuntu:22.04
+
+# Install Python and pip
+RUN apt-get update && apt-get install -y python3 python3-pip
+
+# Optionally install other dependencies, tools, etc.
+
+WORKDIR /app
+COPY . /app
+
+# If you have a requirements.txt, install dependencies
+RUN pip3 install --no-cache-dir -r requirements.txt
+
+CMD [ "tail", "-f", "/dev/null" ]
+
+```
+
+## Step 3: Create devcontainer.json for Docker-Compose
+
+Create **`devcontainer.json`**
+
+**devcontainer.json** (**`./.devcontainer/devcontainer.json`**):
+
+```json
+
+{
+    "name": "TMU Devcontainer",
+    "dockerComposeFile": "docker-compose.yml",
+    "service": "tmu-development",
+    "workspaceFolder": "/app",
+    "extensions": [
+        "ms-python.python",
+    ],
+    "forwardPorts": [],
+    "postCreateCommand": "echo 'Devcontainer is ready'",
+    "remoteUser": "root"
+}
+```
+
+## **Using the Setup**
+
+- After configuring your **`.devcontainer`** directory with the **`Dockerfile`**, **`docker-compose.yml`**, and **`devcontainer.json`**, open your project in VSCode.
+- VSCode may prompt you to reopen the project in a container. If not, you can manually do so by opening the Command Palette (**`F1`** or **`Ctrl+Shift+P`**/**`Cmd+Shift+P`**) and selecting "Remote-Containers: Reopen in Container".
+- This will build your Docker container as defined, including the necessary GPU assignments for CUDA development.
+
+# **3. Development Using Devcontainers on a Remote Machine (SSH)**
+
+Running your development environment on a remote machine can provide significant performance benefits, especially for resource-intensive tasks. This setup requires a bit more initial configuration but has substantial advantages of powerful remote resources and a consistent development environment, like the DGX-2 machines.
+
+## **Prerequisites**
+
+- The remote machine must have Docker installed and running. (This is typically already done unless its your own remote machine)
+- SSH access to the remote machine is set up (refer to the SSH setup guide provided earlier).
+- VSCode and the Remote Development extension pack are installed on your local machine.
+
+## **Setting Up Your Remote Devcontainer Environment**
+
+1. **Connect to Your Remote Machine via SSH**: Open VSCode, then open the Command Palette and select "Remote-SSH: Connect to Host...". Choose your remote machine from the list or add a new SSH connection.
+2. **Initialize Your Project on the Remote Machine**: You can clone your repository or access your project files on the remote machine. This might involve using Git commands within the terminal in VSCode once connected to the remote machine.
+3. **Configure the Devcontainer**: Similar to the local setup, create a **`.devcontainer`** directory in your project on the remote machine with a **`Dockerfile`** and **`devcontainer.json`**. These files might already exist if you cloned a repository already configured for Devcontainer development.
+4. **Open Your Project in a Container Over SSH**: With the remote SSH connection active and your project open in VSCode, use the Command Palette to select "Remote-Containers: Reopen in Container". This will build and start the container on the remote machine, with VSCode connecting to it over SSH.
+5. **Start Developing Remotely**: You can now develop directly on the remote machine, utilizing its resources while benefiting from a consistent, containerized environment controlled by your Devcontainer configuration.
\ No newline at end of file
diff --git a/examples/type_iii_feedback/utils.py b/examples/__init__.py
similarity index 100%
rename from examples/type_iii_feedback/utils.py
rename to examples/__init__.py
diff --git a/examples/autoencoder/DimensionalityReductionDemo.py b/examples/autoencoder/DimensionalityReductionDemo.py
index 04bc356c..0a72dc2a 100644
--- a/examples/autoencoder/DimensionalityReductionDemo.py
+++ b/examples/autoencoder/DimensionalityReductionDemo.py
@@ -1,82 +1,124 @@
+import argparse
+import logging
 import numpy as np
-from time import time
-from sklearn.metrics import f1_score
 from sklearn.metrics.pairwise import cosine_similarity
-
-from scipy.sparse import csr_matrix
-
 from tmu.models.autoencoder.autoencoder import TMAutoEncoder
-
-number_of_features = 12
-
-noise = 0.1
-
-clause_weight_threshold = 0
-
-number_of_examples = 2500
-accumulation = 1
-
-clauses = 10
-T = 8*10
-s = 1.5
-
-print("Number of clauses:", clauses)
-
-output_active = np.array([0, 1], dtype=np.uint32)
-
-X_train = np.random.randint(0, 2, size=(5000, number_of_features), dtype=np.uint32)
-X_train[:,0] = np.where(np.random.rand(5000) <= noise, 1- X_train[:,1],  X_train[:,1])
-
-tm = TMAutoEncoder(clauses, T, s, output_active, max_included_literals=3, accumulation=accumulation, feature_negation=True, platform='CPU', output_balancing=True)
-
-print("\nAccuracy Over 40 Epochs:")
-for e in range(40):
-	start_training = time()
-	tm.fit(X_train, number_of_examples=number_of_examples)
-	stop_training = time()
-
-	print("\nEpoch #%d\n" % (e+1))
-
-	print("Calculating precision\n")
-	precision = []
-	for i in range(output_active.shape[0]):
-		precision.append(tm.clause_precision(i, True, X_train, number_of_examples=500))
-
-	print("Calculating recall\n")
-	recall = []
-	for i in range(output_active.shape[0]):
-		recall.append(tm.clause_recall(i, True, X_train, number_of_examples=500))
-
-	print("Clauses\n")
-
-	for j in range(clauses):
-		print("Clause #%d " % (j), end=' ')
-		for i in range(output_active.shape[0]):
-			print("x%d:W%d:P%.2f:R%.2f " % (i, tm.get_weight(i, j), precision[i][j], recall[i][j]), end=' ')
-
-		l = []
-		for k in range(tm.clause_bank.number_of_literals):
-			if tm.get_ta_action(j, k) == 1:
-				if k < tm.clause_bank.number_of_features:
-					l.append("x%d(%d)" % (k, tm.clause_bank.get_ta_state(j, k)))
-				else:
-					l.append("¬x%d(%d)" % (k-tm.clause_bank.number_of_features, tm.clause_bank.get_ta_state(j, k)))
-		print(" ∧ ".join(l))
-
-	profile = np.empty((output_active.shape[0], clauses))
-	for i in range(output_active.shape[0]):
-		weights = tm.get_weights(i)
-		profile[i,:] = np.where(weights >= clause_weight_threshold, weights, 0)
-
-	similarity = cosine_similarity(profile)
-
-	print("\nWord Similarity\n")
-
-	for i in range(output_active.shape[0]):
-		print("x%d" % (output_active[i]), end=': ')
-		sorted_index = np.argsort(-1*similarity[i,:])
-		for j in range(1, output_active.shape[0]):
-			print("x%d(%.2f) " % (output_active[sorted_index[j]], similarity[i,sorted_index[j]]), end=' ')
-		print()
-
-	print("\nTraining Time: %.2f" % (stop_training - start_training))
+from tmu.tools import BenchmarkTimer
+
+_LOGGER = logging.getLogger(__name__)
+
+def metrics(args):
+    return dict(
+        cosine_similarity=[],
+        train_time=[],
+        recall=[],
+        precision=[],
+        args=vars(args)
+    )
+
+def main(args):
+    experiment_results = metrics(args)
+
+    print("Number of clauses:", args.num_clauses)
+
+    output_active = np.array([0, 1], dtype=np.uint32)
+
+    X_train = np.random.randint(0, 2, size=(5000, args.number_of_features), dtype=np.uint32)
+    X_train[:, 0] = np.where(np.random.rand(5000) <= args.noise, 1 - X_train[:, 1], X_train[:, 1])
+
+    tm = TMAutoEncoder(
+        args.num_clauses,
+        args.T,
+        args.s,
+        output_active,
+        max_included_literals=args.max_included_literals,
+        accumulation=args.accumulation,
+        feature_negation=args.feature_negation,
+        platform=args.platform,
+        output_balancing=args.output_balancing
+    )
+
+    print(f"\nAccuracy Over {args.epochs} Epochs:")
+    for e in range(args.epochs):
+
+        benchmark1 = BenchmarkTimer(logger=_LOGGER, text="Training Time")
+        with benchmark1:
+            tm.fit(X_train, number_of_examples=args.number_of_examples)
+        experiment_results["train_time"].append(benchmark1.elapsed())
+
+        print("\nEpoch #%d\n" % (e + 1))
+
+        print("Calculating precision\n")
+        precision = []
+        for i in range(output_active.shape[0]):
+            precision.append(tm.clause_precision(i, True, X_train, number_of_examples=500))
+        experiment_results["precision"].append(precision)
+
+        print("Calculating recall\n")
+        recall = []
+        for i in range(output_active.shape[0]):
+            recall.append(tm.clause_recall(i, True, X_train, number_of_examples=500))
+        experiment_results["recall"].append(precision)
+        print("Clauses\n")
+
+        for j in range(args.num_clauses):
+            print("Clause #%d " % (j), end=' ')
+            for i in range(output_active.shape[0]):
+                print("x%d:W%d:P%.2f:R%.2f " % (i, tm.get_weight(i, j), precision[i][j], recall[i][j]), end=' ')
+
+            l = []
+            for k in range(tm.clause_bank.number_of_literals):
+                if tm.get_ta_action(j, k) == 1:
+                    if k < tm.clause_bank.number_of_features:
+                        l.append("x%d(%d)" % (k, tm.clause_bank.get_ta_state(j, k)))
+                    else:
+                        l.append(
+                            "¬x%d(%d)" % (k - tm.clause_bank.number_of_features, tm.clause_bank.get_ta_state(j, k)))
+            print(" ∧ ".join(l))
+
+        profile = np.empty((output_active.shape[0], args.num_clauses))
+        for i in range(output_active.shape[0]):
+            weights = tm.get_weights(i)
+            profile[i, :] = np.where(weights >= args.clause_weight_threshold, weights, 0)
+
+        similarity = cosine_similarity(profile)
+        experiment_results["cosine_similarity"].append(list(similarity.flatten()))
+
+        print("\nWord Similarity\n")
+
+        for i in range(output_active.shape[0]):
+            print("x%d" % (output_active[i]), end=': ')
+            sorted_index = np.argsort(-1 * similarity[i, :])
+            for j in range(1, output_active.shape[0]):
+                print("x%d(%.2f) " % (output_active[sorted_index[j]], similarity[i, sorted_index[j]]), end=' ')
+            print()
+
+        print("\nTraining Time: %.2f" % (benchmark1.elapsed()))
+        return experiment_results
+
+
+def default_args(**kwargs):
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--number-of-features", default=12, type=int)
+    parser.add_argument("--noise", default=0.1, type=float)
+    parser.add_argument("--clause-weight-threshold", default=0, type=int)
+    parser.add_argument("--number-of-examples", default=2500, type=int)
+    parser.add_argument("--accumulation", default=1, type=int)
+    parser.add_argument("--num-clauses", default=10, type=bool)
+    parser.add_argument("--T", default=8 * 10, type=int)
+    parser.add_argument("--s", default=1.5, type=float)
+    parser.add_argument("--epochs", default=10, type=int)
+    parser.add_argument("--platform", default="CPU", type=str)
+    parser.add_argument("--output-balancing", default=0.5, type=float)
+    parser.add_argument("--max-included-literals", default=3, type=int)
+    parser.add_argument("--feature-negation", default=True, type=bool)
+    args = parser.parse_args()
+    for key, value in kwargs.items():
+        if key in args.__dict__:
+            setattr(args, key, value)
+    return args
+
+
+if __name__ == "__main__":
+    result = main(default_args())
+    _LOGGER.info(result)
diff --git a/examples/autoencoder/IMDbWordEmbeddingDemo.py b/examples/autoencoder/IMDbWordEmbeddingDemo.py
index f14027c4..44beb46c 100644
--- a/examples/autoencoder/IMDbWordEmbeddingDemo.py
+++ b/examples/autoencoder/IMDbWordEmbeddingDemo.py
@@ -1,5 +1,3 @@
-from time import time
-
 import numpy as np
 from sklearn.feature_extraction.text import CountVectorizer
 from sklearn.metrics.pairwise import cosine_similarity
@@ -9,164 +7,194 @@
 from tmu.data import IMDB
 from tmu.tools import BenchmarkTimer
 import logging
-from tmu.util.cuda_profiler import CudaProfiler
 import os
 import argparse
 
-
 os.environ["CUDA_PROFILE"] = "1"
 
 _LOGGER = logging.getLogger(__name__)
 
 target_words = [
-	'awful',
-	'terrible',
-	'lousy',
-	'abysmal',
-	'crap',
-	'outstanding',
-	'brilliant',
-	'excellent',
-	'superb',
-	'magnificent',
-	'marvellous',
-	'truck',
-	'plane',
-	'car',
-	'cars',
-	'motorcycle',
-	'scary',
-	'frightening',
-	'terrifying',
-	'horrifying',
-	'funny',
-	'comic',
-	'hilarious',
-	'witty'
+    'awful',
+    'terrible',
+    'lousy',
+    'abysmal',
+    'crap',
+    'outstanding',
+    'brilliant',
+    'excellent',
+    'superb',
+    'magnificent',
+    'marvellous',
+    'truck',
+    'plane',
+    'car',
+    'cars',
+    'motorcycle',
+    'scary',
+    'frightening',
+    'terrifying',
+    'horrifying',
+    'funny',
+    'comic',
+    'hilarious',
+    'witty'
 ]
 
-
-
-def main():
-	parser = argparse.ArgumentParser(description='Process some integers.')
-	parser.add_argument('--clause_weight_threshold', type=int, default=0, help='Clause weight threshold')
-	parser.add_argument('--number_of_examples', type=int, default=2500, help='Number of examples')
-	parser.add_argument('--accumulation', type=int, default=25, help='Accumulation')
-	parser.add_argument('--factor', type=int, default=4, help='Factor')
-	parser.add_argument('--clauses', type=int, default=80, help='Clauses')  # factor*20
-	parser.add_argument('--T', type=int, default=160, help='T')  # factor*40
-	parser.add_argument('--s', type=float, default=5.0, help='S')
-	parser.add_argument('--NUM_WORDS', type=int, default=10000, help='Number of words')
-	parser.add_argument('--INDEX_FROM', type=int, default=2, help='Index from')
-	parser.add_argument('--epochs', type=int, default=250, help='Number of epochs')
-	parser.add_argument('--device', type=str, default="CUDA", help='which device to use')
-	parser.add_argument("--output_balancing", type=float, default=0.5, help="output balancing")
-	parser.add_argument("--number_of_output_words", type=int, default=20, help="how many words")
-	parser.add_argument("--max_included_literals", type=int, default=3, help="max included literals")
-	parser.add_argument("--feature_negation", type=bool, default=True, help="feature negation")
-
-
-
-	args = parser.parse_args()
-
-	# load IMDB dataset
-	dataloader = IMDB(num_words=args.NUM_WORDS, index_from=args.INDEX_FROM)
-	dataset = dataloader.get()
-	train_x = dataset["x_train"]
-	train_y = dataset["y_train"]
-	test_x = dataset["x_test"]
-	test_y = dataset["y_test"]
-
-	# get word to indice and id_to_word mappings
-	word_to_id = dataloader.get_word_index()
-	word_to_id = {k: (v + args.INDEX_FROM) for k, v in word_to_id.items()}
-	word_to_id["<PAD>"] = 0
-	word_to_id["<START>"] = 1
-	word_to_id["<UNK>"] = 2
-	id_to_word = {value: key for key, value in word_to_id.items()}
-
-	_LOGGER.info("Producing bit representation...")
-
-	def produce_documents(dsx, dsy):
-		docs = []
-		for i in range(dsy.shape[0]):
-			terms = [id_to_word[word_id].lower() for word_id in dsx[i]]
-			docs.append(terms)
-		return docs
-
-	training_documents = produce_documents(train_x, train_y)
-	testing_documents = produce_documents(test_x, test_y)
-
-	# Create vectorizer
-	def tokenizer(s):
-		return s
-
-	vectorizer_X = CountVectorizer(
-		tokenizer=tokenizer,
-		token_pattern=None,
-		lowercase=False,
-		binary=True
-	)
-
-	X_train = vectorizer_X.fit_transform(training_documents)
-	feature_names = vectorizer_X.get_feature_names_out()
-
-	output_active = np.empty(len(target_words), dtype=np.uint32)
-	for i in range(len(target_words)):
-		target_word = target_words[i]
-		target_id = vectorizer_X.vocabulary_[target_word]
-		output_active[i] = target_id
-
-	tm = TMAutoEncoder(
-		number_of_clauses=args.clauses,
-		T=args.T,
-		s=args.s,
-		output_active=output_active,
-		max_included_literals=args.max_included_literals,
-		accumulation=args.accumulation,
-		feature_negation=args.feature_negation,
-		platform=args.device,
-		output_balancing=args.output_balancing
-	)
-
-	benchmark_train = BenchmarkTimer()
-	benchmark_test = BenchmarkTimer()
-	benchmark_total = BenchmarkTimer()
-
-	_LOGGER.info(f"Accuracy over {args.epochs} epochs:")
-	for e in range(args.epochs):
-		with benchmark_total:
-			_LOGGER.info(f"Epoch {e+1}")
-
-			with benchmark_train:
-				tm.fit(X_train, number_of_examples=args.number_of_examples)
-
-			with benchmark_test:
-				_LOGGER.info("Calculating precision")
-				precision = [tm.clause_precision(i, True, X_train, number_of_examples=500) for i in tqdm(range(len(target_words)), desc="Precision")]
-
-				_LOGGER.info("Calculating recall")
-				recall = [tm.clause_recall(i, True, X_train, number_of_examples=500) for i in tqdm(range(len(target_words)), desc="Recall")]
-
-				_LOGGER.info("Clauses")
-				for j in range(args.clauses):
-					clause_info = " ".join([f"{target_words[i]}:W{tm.get_weight(i, j)}:P{precision[i][j]:.2f}:R{recall[i][j]:.2f}" for i in range(len(target_words))])
-					literals = ["{}{}({})".format("¬" if k >= tm.clause_bank.number_of_features else "", feature_names[k % tm.clause_bank.number_of_features], tm.clause_bank.get_ta_state(j, k)) for k in range(tm.clause_bank.number_of_literals) if tm.get_ta_action(j, k) == 1]
-					_LOGGER.info(f"Clause #{j} {clause_info} {' ∧ '.join(literals)}")
-
-				profile = np.array([np.where(tm.get_weights(i) >= args.clause_weight_threshold, tm.get_weights(i), 0) for i in range(len(target_words))])
-				similarity = cosine_similarity(profile)
-
-				_LOGGER.info("Word Similarity:")
-				for i in range(len(target_words)):
-					sorted_index = np.argsort(-similarity[i, :])
-					similarity_info = " ".join([f"{target_words[sorted_index[j]]}({similarity[i, sorted_index[j]]:.2f})" for j in range(1, len(target_words))])
-					_LOGGER.info(f"{target_words[i]}: {similarity_info}")
-		_LOGGER.info(f"Total time: {benchmark_total.elapsed():.2f}s")
-		_LOGGER.info(f"Training time: {benchmark_train.elapsed():.2f}s")
-		_LOGGER.info(f"Testing time: {benchmark_test.elapsed():.2f}s")
-
+def metrics(args):
+    return dict(
+        accuracy=[],
+        train_time=[],
+        test_time=[],
+        precision=[],
+        recall=[],
+        args=vars(args)
+    )
+
+def main(args):
+    experiment_results = metrics(args)
+
+    # load IMDB dataset
+    dataloader = IMDB(num_words=args.NUM_WORDS, index_from=args.INDEX_FROM)
+    dataset = dataloader.get()
+    train_x = dataset["x_train"]
+    train_y = dataset["y_train"]
+    test_x = dataset["x_test"]
+    test_y = dataset["y_test"]
+
+    # get word to indice and id_to_word mappings
+    word_to_id = dataloader.get_word_index()
+    word_to_id = {k: (v + args.INDEX_FROM) for k, v in word_to_id.items()}
+    word_to_id["<PAD>"] = 0
+    word_to_id["<START>"] = 1
+    word_to_id["<UNK>"] = 2
+    id_to_word = {value: key for key, value in word_to_id.items()}
+
+    _LOGGER.info("Producing bit representation...")
+
+    def produce_documents(dsx, dsy):
+        docs = []
+        for i in range(dsy.shape[0]):
+            terms = [id_to_word[word_id].lower() for word_id in dsx[i]]
+            docs.append(terms)
+        return docs
+
+    training_documents = produce_documents(train_x, train_y)
+    testing_documents = produce_documents(test_x, test_y)
+
+    # Create vectorizer
+    def tokenizer(s):
+        return s
+
+    vectorizer_X = CountVectorizer(
+        tokenizer=tokenizer,
+        token_pattern=None,
+        lowercase=False,
+        binary=True
+    )
+
+    X_train = vectorizer_X.fit_transform(training_documents)
+    feature_names = vectorizer_X.get_feature_names_out()
+
+    output_active = np.empty(len(target_words), dtype=np.uint32)
+    for i in range(len(target_words)):
+        target_word = target_words[i]
+        target_id = vectorizer_X.vocabulary_[target_word]
+        output_active[i] = target_id
+
+    tm = TMAutoEncoder(
+        number_of_clauses=args.clauses,
+        T=args.T,
+        s=args.s,
+        output_active=output_active,
+        max_included_literals=args.max_included_literals,
+        accumulation=args.accumulation,
+        feature_negation=args.feature_negation,
+        platform=args.device,
+        output_balancing=args.output_balancing
+    )
+
+    benchmark_train = BenchmarkTimer()
+    benchmark_test = BenchmarkTimer()
+    benchmark_total = BenchmarkTimer()
+
+    _LOGGER.info(f"Accuracy over {args.epochs} epochs:")
+    for e in range(args.epochs):
+        with benchmark_total:
+            _LOGGER.info(f"Epoch {e + 1}")
+
+            with benchmark_train:
+                tm.fit(X_train, number_of_examples=args.number_of_examples)
+            experiment_results["train_time"].append(benchmark_train.elapsed())
+
+            with benchmark_test:
+                _LOGGER.info("Calculating precision")
+                precision = [tm.clause_precision(i, True, X_train, number_of_examples=500) for i in
+                             tqdm(range(len(target_words)), desc="Precision")]
+
+                _LOGGER.info("Calculating recall")
+                recall = [tm.clause_recall(i, True, X_train, number_of_examples=500) for i in
+                          tqdm(range(len(target_words)), desc="Recall")]
+
+                experiment_results["precision"].append(precision)
+                experiment_results["recall"].append(precision)
+
+                _LOGGER.info("Clauses")
+                for j in range(args.clauses):
+                    clause_info = " ".join(
+                        [f"{target_words[i]}:W{tm.get_weight(i, j)}:P{precision[i][j]:.2f}:R{recall[i][j]:.2f}" for i in
+                         range(len(target_words))])
+                    literals = ["{}{}({})".format("¬" if k >= tm.clause_bank.number_of_features else "",
+                                                  feature_names[k % tm.clause_bank.number_of_features],
+                                                  tm.clause_bank.get_ta_state(j, k)) for k in
+                                range(tm.clause_bank.number_of_literals) if tm.get_ta_action(j, k) == 1]
+                    _LOGGER.info(f"Clause #{j} {clause_info} {' ∧ '.join(literals)}")
+
+                profile = np.array(
+                    [np.where(tm.get_weights(i) >= args.clause_weight_threshold, tm.get_weights(i), 0) for i in
+                     range(len(target_words))])
+                similarity = cosine_similarity(profile)
+
+                _LOGGER.info("Word Similarity:")
+                for i in range(len(target_words)):
+                    sorted_index = np.argsort(-similarity[i, :])
+                    similarity_info = " ".join(
+                        [f"{target_words[sorted_index[j]]}({similarity[i, sorted_index[j]]:.2f})" for j in
+                         range(1, len(target_words))])
+                    _LOGGER.info(f"{target_words[i]}: {similarity_info}")
+
+            experiment_results["test_time"].append(benchmark_train.elapsed())
+
+        _LOGGER.info(f"Total time: {benchmark_total.elapsed():.2f}s")
+        _LOGGER.info(f"Training time: {benchmark_train.elapsed():.2f}s")
+        _LOGGER.info(f"Testing time: {benchmark_test.elapsed():.2f}s")
+    return experiment_results
+
+
+def default_args(**kwargs):
+    parser = argparse.ArgumentParser(description='Process some integers.')
+    parser.add_argument('--clause_weight_threshold', type=int, default=0, help='Clause weight threshold')
+    parser.add_argument('--number_of_examples', type=int, default=2500, help='Number of examples')
+    parser.add_argument('--accumulation', type=int, default=25, help='Accumulation')
+    parser.add_argument('--factor', type=int, default=4, help='Factor')
+    parser.add_argument('--clauses', type=int, default=80, help='Clauses')  # factor*20
+    parser.add_argument('--T', type=int, default=160, help='T')  # factor*40
+    parser.add_argument('--s', type=float, default=5.0, help='S')
+    parser.add_argument('--NUM_WORDS', type=int, default=10000, help='Number of words')
+    parser.add_argument('--INDEX_FROM', type=int, default=2, help='Index from')
+    parser.add_argument('--epochs', type=int, default=250, help='Number of epochs')
+    parser.add_argument('--device', type=str, default="CUDA", help='which device to use')
+    parser.add_argument("--output_balancing", type=float, default=0.5, help="output balancing")
+    parser.add_argument("--number_of_output_words", type=int, default=20, help="how many words")
+    parser.add_argument("--max_included_literals", type=int, default=3, help="max included literals")
+    parser.add_argument("--feature_negation", type=bool, default=True, help="feature negation")
+    args = parser.parse_args()
+    for key, value in kwargs.items():
+        if key in args.__dict__:
+            setattr(args, key, value)
+    return args
 
 
 if __name__ == "__main__":
-	main()
+    results = main(default_args())
+    _LOGGER.info(results)
diff --git a/examples/classification/CIFAR2Demo3x3LiteralBudget.py b/examples/classification/CIFAR2Demo3x3LiteralBudget.py
index e562f6bd..457745da 100644
--- a/examples/classification/CIFAR2Demo3x3LiteralBudget.py
+++ b/examples/classification/CIFAR2Demo3x3LiteralBudget.py
@@ -1,76 +1,173 @@
+import argparse
+import json
+import logging
+from multiprocessing import get_context
+
+from keras.datasets import cifar10
 from tmu.models.classification.vanilla_classifier import TMClassifier
 import numpy as np
-from time import time
 import ssl
+from tmu.tools import BenchmarkTimer
 
 ssl._create_default_https_context = ssl._create_unverified_context
 
-from keras.datasets import cifar10
-
-animals = np.array([2, 3, 4, 5, 6, 7])
-
-max_included_literals = 32
-clauses = 8000
-T = int(clauses * 0.75)
-s = 10.0
-patch_size = 3
-resolution = 8
-number_of_state_bits_ta = 8
-literal_drop_p = 0.0
-
-epochs = 250
-ensembles = 5
-
-classes = 10
-
-(X_train_org, Y_train), (X_test_org, Y_test) = cifar10.load_data()
-
-Y_train = Y_train.reshape(Y_train.shape[0])
-Y_test = Y_test.reshape(Y_test.shape[0])
-
-X_train = np.empty((X_train_org.shape[0], X_train_org.shape[1], X_train_org.shape[2], X_train_org.shape[3], resolution),
-                   dtype=np.uint8)
-for z in range(resolution):
-    X_train[:, :, :, :, z] = X_train_org[:, :, :, :] >= (z + 1) * 255 / (resolution + 1)
-
-X_test = np.empty((X_test_org.shape[0], X_test_org.shape[1], X_test_org.shape[2], X_test_org.shape[3], resolution),
-                  dtype=np.uint8)
-for z in range(resolution):
-    X_test[:, :, :, :, z] = X_test_org[:, :, :, :] >= (z + 1) * 255 / (resolution + 1)
-
-X_train = X_train.reshape((X_train_org.shape[0], X_train_org.shape[1], X_train_org.shape[2], 3 * resolution))
-X_test = X_test.reshape((X_test_org.shape[0], X_test_org.shape[1], X_test_org.shape[2], 3 * resolution))
-
-Y_train = np.where(np.isin(Y_train, animals), 1, 0)
-Y_test = np.where(np.isin(Y_test, animals), 1, 0)
-
-f = open("cifar2_%.1f_%d_%d_%d_%.2f_%d_%d.txt" % (
-s, clauses, T, patch_size, literal_drop_p, resolution, max_included_literals), "w+")
-for ensemble in range(ensembles):
-    tm = TMClassifier(clauses, T, s, platform='GPU', patch_dim=(patch_size, patch_size),
-                      number_of_state_bits_ta=number_of_state_bits_ta, weighted_clauses=True,
-                      literal_drop_p=literal_drop_p, max_included_literals=max_included_literals)
-    for epoch in range(epochs):
-        start_training = time()
-        tm.fit(X_train, Y_train)
-        stop_training = time()
-
-        start_testing = time()
-        result_test = 100 * (tm.predict(X_test) == Y_test).mean()
-        stop_testing = time()
-
-        result_train = 100 * (tm.predict(X_train) == Y_train).mean()
+_LOGGER = logging.getLogger(__name__)
+
+
+def preprocess_cifar10_data(resolution, animals):
+    """
+    Preprocess CIFAR-10 images and labels.
+
+    Parameters:
+    - resolution: The number of bins to quantize the pixel values into.
+    - animals: An array of CIFAR-10 label indices to be considered as positive samples (1), with all others as negative samples (0).
+
+    Returns:
+    - X_train, Y_train: Processed training images and labels.
+    - X_test, Y_test: Processed testing images and labels.
+    """
+    # Load CIFAR-10 data
+    (X_train_org, Y_train), (X_test_org, Y_test) = cifar10.load_data()
+
+    # Flatten Y arrays
+    Y_train, Y_test = Y_train.reshape(-1), Y_test.reshape(-1)
+
+    # Initialize empty arrays for quantized images
+    X_train = np.empty(
+        (X_train_org.shape[0], X_train_org.shape[1], X_train_org.shape[2], X_train_org.shape[3], resolution),
+        dtype=np.uint8)
+    X_test = np.empty((X_test_org.shape[0], X_test_org.shape[1], X_test_org.shape[2], X_test_org.shape[3], resolution),
+                      dtype=np.uint8)
+
+    # Quantize pixel values
+    for z in range(resolution):
+        threshold = (z + 1) * 255 / (resolution + 1)
+        X_train[:, :, :, :, z] = X_train_org >= threshold
+        X_test[:, :, :, :, z] = X_test_org >= threshold
+
+    # Reshape quantized images and convert to uint32
+    X_train = X_train.reshape(
+        (X_train_org.shape[0], X_train_org.shape[1], X_train_org.shape[2], 3 * resolution)).astype(np.uint32)
+    X_test = X_test.reshape((X_test_org.shape[0], X_test_org.shape[1], X_test_org.shape[2], 3 * resolution)).astype(
+        np.uint32)
+
+    # Process labels: 1 for animals, 0 for others
+    Y_train = np.where(np.isin(Y_train, animals), 1, 0).astype(np.uint32)
+    Y_test = np.where(np.isin(Y_test, animals), 1, 0).astype(np.uint32)
+
+    return X_train, Y_train, X_test, Y_test
+
+
+def metrics(args):
+    return dict(
+        accuracy=[],
+        accuracy_train=[],
+        number_of_includes=[],
+        train_time=[],
+        test_time=[],
+        str_result=[],
+        args=vars(args)
+    )
+
+def run_ensemble(ensemble_params):
+    # Unpack parameters
+    args, X_train, Y_train, X_test, Y_test, ensemble = ensemble_params
+
+    T = int(args.clauses * 0.75)
+    tm = TMClassifier(
+        args.clauses,
+        T,
+        args.s,
+        platform=args.platform,
+        patch_dim=(args.patch_size, args.patch_size),
+        number_of_state_bits_ta=args.number_of_state_bits_ta,
+        weighted_clauses=args.weighted_clauses,
+        literal_drop_p=args.literal_drop_p,
+        max_included_literals=args.max_included_literals
+    )
+
+    ensemble_results = metrics(args)
+
+    for epoch in range(args.epochs):
+        # Benchmark training time
+        benchmark1 = BenchmarkTimer()
+        with benchmark1:
+            tm.fit(X_train, Y_train)
+        ensemble_results["train_time"].append(benchmark1.elapsed())
+
+        # Benchmark testing time and calculate accuracy
+        benchmark2 = BenchmarkTimer()
+        with benchmark2:
+            accuracy_test = 100 * (tm.predict(X_test) == Y_test).mean()
+        ensemble_results["test_time"].append(benchmark2.elapsed())
+        ensemble_results["accuracy"].append(accuracy_test)
+
+        # Calculate additional metrics
+        accuracy_train = 100 * (tm.predict(X_train) == Y_train).mean()
+        ensemble_results["accuracy_train"].append(accuracy_train)
 
         number_of_includes = 0
         for i in range(2):
-            for j in range(clauses):
+            for j in range(args.clauses):
                 number_of_includes += tm.number_of_include_actions(i, j)
-        number_of_includes /= 2 * clauses
-
-        print("%d %d %.2f %.2f %.2f %.2f %.2f" % (
-        ensemble, epoch, number_of_includes, result_train, result_test, stop_training - start_training, stop_testing - start_testing))
-        print("%d %d %.2f %.2f %.2f %.2f %.2f" % (
-        ensemble, epoch, number_of_includes, result_train, result_test, stop_training - start_training, stop_testing - start_testing),
-              file=f)
-        f.flush()
-f.close()
+        number_of_includes /= 2 * args.clauses
+        ensemble_results["number_of_includes"].append(number_of_includes)
+
+        out_str = "%d %d %.2f %.2f %.2f %.2f %.2f" % (
+            ensemble, epoch, number_of_includes, accuracy_train, accuracy_test, benchmark1.elapsed(),
+            benchmark2.elapsed())
+        ensemble_results["str_result"].append(out_str)
+        _LOGGER.info(out_str)
+    return ensemble, ensemble_results
+
+def run_ensemble_no_mp(ensemble_params):
+    """A wrapper to call `run_ensemble` without multiprocessing."""
+    # Just directly call run_ensemble function without any multiprocessing overhead
+    return run_ensemble(ensemble_params)
+
+
+def main(args):
+    animals = np.array([2, 3, 4, 5, 6, 7])
+    X_train, Y_train, X_test, Y_test = preprocess_cifar10_data(args.resolution, animals)
+
+    ensemble_params = [(args, X_train, Y_train, X_test, Y_test, ensemble) for ensemble in range(args.ensembles)]
+
+    # Check if multiprocessing should be used
+    if args.use_multiprocessing:
+        with get_context("spawn").Pool(processes=2) as pool:
+            results = pool.map(run_ensemble, ensemble_params)
+    else:
+        # Run ensembles sequentially without multiprocessing
+        results = list(map(run_ensemble_no_mp, ensemble_params))
+
+    experiment_results = {}
+    for ensemble, ensemble_results in results:
+        experiment_results[f'ensemble_{ensemble}'] = ensemble_results
+
+    return experiment_results
+
+
+def default_args(**kwargs):
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--max_included_literals", type=int, default=32)
+    parser.add_argument("--clauses", type=int, default=8000)
+    parser.add_argument("--s", type=float, default=10.0)
+    parser.add_argument("--platform", type=str, default="GPU")
+    parser.add_argument("--patch_size", type=int, default=3)
+    parser.add_argument("--resolution", type=int, default=8)
+    parser.add_argument("--number_of_state_bits_ta", type=int, default=8)
+    parser.add_argument("--literal_drop_p", type=float, default=0.0)
+    parser.add_argument("--epochs", type=int, default=250)
+    parser.add_argument("--ensembles", type=int, default=5)
+    parser.add_argument("--weighted-clauses", type=bool, default=True)
+    parser.add_argument("--use_multiprocessing", action='store_true', help="Use multiprocessing to run ensembles in parallel")
+    args = parser.parse_args()
+    for key, value in kwargs.items():
+        if key in args.__dict__:
+            setattr(args, key, value)
+    return args
+
+
+if __name__ == "__main__":
+    results = main(default_args())
+    _LOGGER.info(results)
diff --git a/examples/classification/IMDbSparseAbsorbingTextCategorizationDemo.py b/examples/classification/IMDbSparseAbsorbingTextCategorizationDemo.py
index 61025036..36db4f0e 100644
--- a/examples/classification/IMDbSparseAbsorbingTextCategorizationDemo.py
+++ b/examples/classification/IMDbSparseAbsorbingTextCategorizationDemo.py
@@ -4,28 +4,26 @@
 import keras
 from sklearn.feature_selection import SelectKBest
 from sklearn.feature_selection import chi2
-from keras.datasets import imdb
 from sklearn.feature_extraction.text import CountVectorizer
 from tmu.models.classification.vanilla_classifier import TMClassifier
-from time import time
+
+from tmu.tools import BenchmarkTimer
 
 _LOGGER = logging.getLogger(__name__)
 
-if __name__ == "__main__":
+def metrics(args):
+    return dict(
+        accuracy=[],
+        absorbed=[],
+        unallocated=[],
+        train_time=[],
+        test_time=[],
+        args=vars(args)
+    )
 
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--num_clauses", default=10000, type=int)
-    parser.add_argument("--T", default=8000, type=int)
-    parser.add_argument("--s", default=2.0, type=float)
-    parser.add_argument("--device", default="GPU", type=str)
-    parser.add_argument("--weighted_clauses", default=True, type=bool)
-    parser.add_argument("--epochs", default=40, type=int)
-    parser.add_argument("--clause_drop_p", default=0.75, type=float)
-    parser.add_argument("--max-ngram", default=2, type=int)
-    parser.add_argument("--features", default=5000, type=int)
-    parser.add_argument("--imdb-num-words", default=5000, type=int)
-    parser.add_argument("--imdb-index-from", default=2, type=int)
-    args = parser.parse_args()
+
+def main(args):
+    experiment_results = metrics(args)
 
     _LOGGER.info("Preparing dataset")
     train, test = keras.datasets.imdb.load_data(num_words=args.imdb_num_words, index_from=args.imdb_index_from)
@@ -80,28 +78,82 @@
     SKB.fit(X_train, Y_train)
 
     selected_features = SKB.get_support(indices=True)
-    X_train = SKB.transform(X_train)
-    X_test = SKB.transform(X_test)
+    X_train = SKB.transform(X_train).astype(np.uint32)
+    X_test = SKB.transform(X_test).astype(np.uint32)
 
     _LOGGER.info("Selecting Features.... Done!")
 
-    tm = TMClassifier(args.num_clauses, args.T, args.s, platform='CPU_sparse', weighted_clauses=args.weighted_clauses, clause_drop_p=args.clause_drop_p, literal_insertion_state=100, literal_sampling=0.1, absorbing=90)
+    tm = TMClassifier(
+        args.num_clauses,
+        args.T,
+        args.s,
+        platform=args.platform,
+        weighted_clauses=args.weighted_clauses,
+        clause_drop_p=args.clause_drop_p,
+        literal_insertion_state=args.literal_insertion_state,
+        literal_sampling=args.literal_sampling,
+        absorbing=args.absorbing
+    )
+
+    for e in range(args.epochs):
 
-    for e in range(60):
-        start_training = time()
-        tm.fit(X_train, Y_train)
-        stop_training = time()
+        benchmark1 = BenchmarkTimer(logger=_LOGGER, text="Training Time")
+        with benchmark1:
+            tm.fit(X_train, Y_train)
+        experiment_results["train_time"].append(benchmark1.elapsed())
 
-        start_testing = time()
-        result = 100*(tm.predict(X_test) == Y_test).mean()
-        stop_testing = time()
+        benchmark2 = BenchmarkTimer(logger=_LOGGER, text="Testing Time")
+        with benchmark2:
+            result = 100 * (tm.predict(X_test) == Y_test).mean()
+        experiment_results["test_time"].append(benchmark2.elapsed())
+        experiment_results["accuracy"].append(int(result))
 
         absorbed = 0.0
         unallocated = 0
         for i in range(2):
-                for j in range(args.num_clauses):
-                        absorbed += 1.0 - (tm.number_of_include_actions(i, j) + tm.number_of_exclude_actions(i, j)) / (X_train.shape[1]*2)
-                        unallocated += tm.number_of_unallocated_literals(i, j)
-        absorbed = 100 * absorbed / (2*args.num_clauses)
+            for j in range(args.num_clauses):
+                absorbed += 1.0 - (tm.number_of_include_actions(i, j) + tm.number_of_exclude_actions(i, j)) / (
+                            X_train.shape[1] * 2)
+                unallocated += tm.number_of_unallocated_literals(i, j)
+        absorbed = 100 * absorbed / (2 * args.num_clauses)
+
+        experiment_results["absorbed"].append(int(absorbed))
+        experiment_results["unallocated"].append(int(unallocated))
+
+        _LOGGER.info("#%d Accuracy: %.2f%% Absorbed: %.2f%% Unallocated: %d Training: %.2fs Testing: %.2fs" % (
+            e + 1,
+            result,
+            absorbed,
+            unallocated,
+            benchmark1.elapsed(),
+            benchmark2.elapsed()
+        ))
+
+    return experiment_results
+
+def default_args(**kwargs):
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--num_clauses", default=10000, type=int)
+    parser.add_argument("--T", default=8000, type=int)
+    parser.add_argument("--s", default=2.0, type=float)
+    parser.add_argument("--platform", default="CPU_sparse", type=str)
+    parser.add_argument("--weighted_clauses", default=True, type=bool)
+    parser.add_argument("--epochs", default=60, type=int)
+    parser.add_argument("--clause_drop_p", default=0.75, type=float)
+    parser.add_argument("--max-ngram", default=2, type=int)
+    parser.add_argument("--features", default=5000, type=int)
+    parser.add_argument("--imdb-num-words", default=5000, type=int)
+    parser.add_argument("--imdb-index-from", default=2, type=int)
+    parser.add_argument("--literal-insertion-state", default=100, type=int)
+    parser.add_argument("--literal-sampling", default=0.1, type=float)
+    parser.add_argument("--absorbing", default=90, type=int)
+    args = parser.parse_args()
+    for key, value in kwargs.items():
+        if key in args.__dict__:
+            setattr(args, key, value)
+    return args
+
+if __name__ == "__main__":
 
-        print("#%d Accuracy: %.2f%% Absorbed: %.2f%% Unallocated: %d Training: %.2fs Testing: %.2fs" % (e+1, result, absorbed, unallocated, stop_training-start_training, stop_testing-start_testing))
\ No newline at end of file
+    results = main(default_args())
+    _LOGGER.info(results)
diff --git a/examples/classification/IMDbTextCategorizationDemo.py b/examples/classification/IMDbTextCategorizationDemo.py
index 46d82bc3..04964c77 100644
--- a/examples/classification/IMDbTextCategorizationDemo.py
+++ b/examples/classification/IMDbTextCategorizationDemo.py
@@ -4,28 +4,22 @@
 import keras
 from sklearn.feature_selection import SelectKBest
 from sklearn.feature_selection import chi2
-from keras.datasets import imdb
 from sklearn.feature_extraction.text import CountVectorizer
 from tmu.models.classification.vanilla_classifier import TMClassifier
 from tmu.tools import BenchmarkTimer
 
 _LOGGER = logging.getLogger(__name__)
 
-if __name__ == "__main__":
+def metrics(args):
+    return dict(
+        accuracy=[],
+        train_time=[],
+        test_time=[],
+        args=vars(args)
+    )
 
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--num_clauses", default=10000, type=int)
-    parser.add_argument("--T", default=8000, type=int)
-    parser.add_argument("--s", default=2.0, type=float)
-    parser.add_argument("--device", default="GPU", type=str)
-    parser.add_argument("--weighted_clauses", default=True, type=bool)
-    parser.add_argument("--epochs", default=40, type=int)
-    parser.add_argument("--clause_drop_p", default=0.75, type=float)
-    parser.add_argument("--max-ngram", default=2, type=int)
-    parser.add_argument("--features", default=5000, type=int)
-    parser.add_argument("--imdb-num-words", default=5000, type=int)
-    parser.add_argument("--imdb-index-from", default=2, type=int)
-    args = parser.parse_args()
+def main(args):
+    experiment_results = metrics(args)
 
     _LOGGER.info("Preparing dataset")
     train, test = keras.datasets.imdb.load_data(num_words=args.imdb_num_words, index_from=args.imdb_index_from)
@@ -80,23 +74,57 @@
     SKB.fit(X_train, Y_train)
 
     selected_features = SKB.get_support(indices=True)
-    X_train = SKB.transform(X_train).toarray()
-    X_test = SKB.transform(X_test).toarray()
+    X_train = SKB.transform(X_train).toarray().astype(np.uint32)
+    X_test = SKB.transform(X_test).toarray().astype(np.uint32)
 
     _LOGGER.info("Selecting Features.... Done!")
 
-    tm = TMClassifier(args.num_clauses, args.T, args.s, platform=args.device, weighted_clauses=args.weighted_clauses,
-                      clause_drop_p=args.clause_drop_p)
+    tm = TMClassifier(
+        args.num_clauses,
+        args.T, args.s,
+        platform=args.platform,
+        weighted_clauses=args.weighted_clauses,
+        clause_drop_p=args.clause_drop_p
+    )
 
     _LOGGER.info(f"Running {TMClassifier} for {args.epochs}")
     for epoch in range(args.epochs):
         benchmark1 = BenchmarkTimer(logger=_LOGGER, text="Training Time")
         with benchmark1:
             tm.fit(X_train, Y_train)
+        experiment_results["train_time"].append(benchmark1.elapsed())
 
         benchmark2 = BenchmarkTimer(logger=_LOGGER, text="Testing Time")
         with benchmark2:
             result = 100 * (tm.predict(X_test) == Y_test).mean()
+        experiment_results["test_time"].append(benchmark2.elapsed())
+        experiment_results["accuracy"].append(result)
 
         _LOGGER.info(f"Epoch: {epoch + 1}, Accuracy: {result:.2f}, Training Time: {benchmark1.elapsed():.2f}s, "
                      f"Testing Time: {benchmark2.elapsed():.2f}s")
+
+    return experiment_results
+
+def default_args(**kwargs):
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--num_clauses", default=10000, type=int)
+    parser.add_argument("--T", default=8000, type=int)
+    parser.add_argument("--s", default=2.0, type=float)
+    parser.add_argument("--platform", default="GPU", type=str)
+    parser.add_argument("--weighted_clauses", default=True, type=bool)
+    parser.add_argument("--epochs", default=40, type=int)
+    parser.add_argument("--clause_drop_p", default=0.75, type=float)
+    parser.add_argument("--max-ngram", default=2, type=int)
+    parser.add_argument("--features", default=5000, type=int)
+    parser.add_argument("--imdb-num-words", default=5000, type=int)
+    parser.add_argument("--imdb-index-from", default=2, type=int)
+    args = parser.parse_args()
+    for key, value in kwargs.items():
+        if key in args.__dict__:
+            setattr(args, key, value)
+    return args
+
+if __name__ == "__main__":
+
+    results = main(default_args())
+    _LOGGER.info(results)
diff --git a/examples/classification/InterpretabilityDemo.py b/examples/classification/InterpretabilityDemo.py
index 044c2ca6..c4de4fa1 100644
--- a/examples/classification/InterpretabilityDemo.py
+++ b/examples/classification/InterpretabilityDemo.py
@@ -1,92 +1,143 @@
+import argparse
+import logging
+
 from tmu.models.classification.vanilla_classifier import TMClassifier
 import numpy as np
-from time import time
-
-number_of_features = 20
-noise = 0.1
-
-X_train = np.random.randint(0, 2, size=(5000, number_of_features), dtype=np.uint32)
-Y_train = np.logical_xor(X_train[:,0], X_train[:,1]).astype(dtype=np.uint32)
-Y_train = np.where(np.random.rand(5000) <= noise, 1-Y_train, Y_train) # Adds noise
-
-X_test = np.random.randint(0, 2, size=(5000, number_of_features), dtype=np.uint32)
-Y_test = np.logical_xor(X_test[:,0], X_test[:,1]).astype(dtype=np.uint32)
-
-tm = TMClassifier(10, 15, 3.0, platform='CPU', boost_true_positive_feedback=0)
-
-for i in range(20):
-	tm.fit(X_train, Y_train)
-
-print("Accuracy:", 100*(tm.predict(X_test) == Y_test).mean())
-
-np.set_printoptions(threshold=np.inf, linewidth=200, precision=2, suppress=True)
-
-print("\nClass 0 Positive Clauses:\n")
-
-precision = tm.clause_precision(0, 0, X_test, Y_test)
-recall = tm.clause_recall(0, 0, X_test, Y_test)
-
-for j in range(5):
-	print("Clause #%d W:%d P:%.2f R:%.2f " % (j, tm.get_weight(0, 0, j), precision[j], recall[j]), end=' ')
-	l = []
-	for k in range(number_of_features*2):
-		if tm.get_ta_action(j, k, the_class = 0, polarity = 0):
-			if k < number_of_features:
-				l.append(" x%d" % (k))
-			else:
-				l.append("¬x%d" % (k-number_of_features))
-	print(" ∧ ".join(l))
-
-print("\nClass 0 Negative Clauses:\n")
-
-precision = tm.clause_precision(0, 1, X_test, Y_test)
-recall = tm.clause_recall(0, 1, X_test, Y_test)
-
-for j in range(5):
-	print("Clause #%d W:%d P:%.2f R:%.2f " % (j, tm.get_weight(0, 1, j), precision[j], recall[j]), end=' ')
-	l = []
-	for k in range(number_of_features*2):
-		if tm.get_ta_action(j, k, the_class = 0, polarity = 1):
-			if k < number_of_features:
-				l.append(" x%d" % (k))
-			else:
-				l.append("¬x%d" % (k-number_of_features))
-	print(" ∧ ".join(l))
-
-print("\nClass 1 Positive Clauses:\n")
-
-precision = tm.clause_precision(1, 0, X_test, Y_test)
-recall = tm.clause_recall(1, 0, X_test, Y_test)
-
-for j in range(5):
-	print("Clause #%d W:%d P:%.2f R:%.2f " % (j, tm.get_weight(1, 0, j), precision[j], recall[j]), end=' ')
-	l = []
-	for k in range(number_of_features*2):
-		if tm.get_ta_action(j, k, the_class = 1, polarity = 0):
-			if k < number_of_features:
-				l.append(" x%d" % (k))
-			else:
-				l.append("¬x%d" % (k-number_of_features))
-	print(" ∧ ".join(l))
-
-print("\nClass 1 Negative Clauses:\n")
-
-precision = tm.clause_precision(1, 1, X_test, Y_test)
-recall = tm.clause_recall(1, 1, X_test, Y_test)
-
-for j in range(5):
-	print("Clause #%d W:%d P:%.2f R:%.2f " % (j, tm.get_weight(1, 1, j), precision[j], recall[j]), end=' ')
-	l = []
-	for k in range(number_of_features*2):
-		if tm.get_ta_action(j, k, the_class = 1, polarity = 1):
-			if k < number_of_features:
-				l.append(" x%d" % (k))
-			else:
-				l.append("¬x%d" % (k-number_of_features))
-	print(" ∧ ".join(l))
-
-print("\nClause Co-Occurence Matrix:\n")
-print(tm.clause_co_occurrence(X_test, percentage=True).toarray())
-
-print("\nLiteral Frequency:\n")
-print(tm.literal_clause_frequency())
+
+_LOGGER = logging.getLogger(__name__)
+
+def metrics(args):
+    return dict(
+        accuracy=[],
+        class_0_precision_positive=[],
+        class_0_recall_positive=[],
+        class_0_recall_negative=[],
+        class_0_precision_negative=[],
+        class_1_precision_positive=[],
+        class_1_recall_positive=[],
+        class_1_recall_negative=[],
+        class_1_precision_negative=[],
+        literal_frequency=None,
+        args=vars(args)
+    )
+
+def main(args):
+    experiment_results = metrics(args)
+
+    X_train = np.random.randint(0, 2, size=(5000, args.number_of_features), dtype=np.uint32)
+    Y_train = np.logical_xor(X_train[:, 0], X_train[:, 1]).astype(dtype=np.uint32)
+    Y_train = np.where(np.random.rand(5000) <= args.noise, 1 - Y_train, Y_train)  # Adds noise
+
+    X_test = np.random.randint(0, 2, size=(5000, args.number_of_features), dtype=np.uint32)
+    Y_test = np.logical_xor(X_test[:, 0], X_test[:, 1]).astype(dtype=np.uint32)
+
+    tm = TMClassifier(args.number_of_clauses, args.T, args.s, platform=args.platform, boost_true_positive_feedback=0)
+
+    for i in range(20):
+        tm.fit(X_train, Y_train)
+        accuracy = 100 * (tm.predict(X_test) == Y_test).mean()
+        experiment_results["accuracy"].append(accuracy)
+        print("Accuracy:", accuracy)
+
+    np.set_printoptions(threshold=np.inf, linewidth=200, precision=2, suppress=True)
+
+    print("\nClass 0 Positive Clauses:\n")
+    precision = tm.clause_precision(0, 0, X_test, Y_test)
+    recall = tm.clause_recall(0, 0, X_test, Y_test)
+    experiment_results["class_0_precision_positive"].append(list(np.asarray(precision)))
+    experiment_results["class_0_recall_positive"].append(list(np.asarray(recall)))
+
+    for j in range(args.number_of_clauses // 2):
+        print("Clause #%d W:%d P:%.2f R:%.2f " % (j, tm.get_weight(0, 0, j), precision[j], recall[j]), end=' ')
+        l = []
+        for k in range(args.number_of_features * 2):
+            if tm.get_ta_action(j, k, the_class=0, polarity=0):
+                if k < args.number_of_features:
+                    l.append(" x%d(%d)" % (k, tm.get_ta_state(j, k, the_class=0, polarity=0)))
+                else:
+                    l.append("¬x%d(%d)" % (k - args.number_of_features, tm.get_ta_state(j, k, the_class=0, polarity=0)))
+        print(" ∧ ".join(l))
+
+    print("\nClass 0 Negative Clauses:\n")
+
+    precision = tm.clause_precision(0, 1, X_test, Y_test)
+    recall = tm.clause_recall(0, 1, X_test, Y_test)
+    experiment_results["class_0_precision_negative"].append(list(np.asarray(precision)))
+    experiment_results["class_0_recall_negative"].append(list(np.asarray(recall)))
+
+    for j in range(args.number_of_clauses // 2):
+        print("Clause #%d W:%d P:%.2f R:%.2f " % (j, tm.get_weight(0, 1, j), precision[j], recall[j]), end=' ')
+        l = []
+        for k in range(args.number_of_features * 2):
+            if tm.get_ta_action(j, k, the_class=0, polarity=1):
+                if k < args.number_of_features:
+                    l.append(" x%d(%d)" % (k, tm.get_ta_state(j, k, the_class=0, polarity=1)))
+                else:
+                    l.append("¬x%d(%d)" % (k - args.number_of_features, tm.get_ta_state(j, k, the_class=0, polarity=1)))
+        print(" ∧ ".join(l))
+
+    print("\nClass 1 Positive Clauses:\n")
+
+    precision = tm.clause_precision(1, 0, X_test, Y_test)
+    recall = tm.clause_recall(1, 0, X_test, Y_test)
+    experiment_results["class_1_precision_positive"].append(list(np.asarray(precision)))
+    experiment_results["class_1_recall_positive"].append(list(np.asarray(recall)))
+
+    for j in range(args.number_of_clauses // 2):
+        print("Clause #%d W:%d P:%.2f R:%.2f " % (j, tm.get_weight(1, 0, j), precision[j], recall[j]), end=' ')
+        l = []
+        for k in range(args.number_of_features * 2):
+            if tm.get_ta_action(j, k, the_class=1, polarity=0):
+                if k < args.number_of_features:
+                    l.append(" x%d(%d)" % (k, tm.get_ta_state(j, k, the_class=1, polarity=0)))
+                else:
+                    l.append("¬x%d(%d)" % (k - args.number_of_features, tm.get_ta_state(j, k, the_class=1, polarity=0)))
+        print(" ∧ ".join(l))
+
+    print("\nClass 1 Negative Clauses:\n")
+
+    precision = tm.clause_precision(1, 1, X_test, Y_test)
+    recall = tm.clause_recall(1, 1, X_test, Y_test)
+    experiment_results["class_1_precision_negative"].append(list(np.asarray(precision)))
+    experiment_results["class_1_recall_negative"].append(list(np.asarray(recall)))
+
+    for j in range(args.number_of_clauses // 2):
+        print("Clause #%d W:%d P:%.2f R:%.2f " % (j, tm.get_weight(1, 1, j), precision[j], recall[j]), end=' ')
+        l = []
+        for k in range(args.number_of_features * 2):
+            if tm.get_ta_action(j, k, the_class=1, polarity=1):
+                if k < args.number_of_features:
+                    l.append(" x%d(%d)" % (k, tm.get_ta_state(j, k, the_class=1, polarity=1)))
+                else:
+                    l.append("¬x%d(%d)" % (k - args.number_of_features, tm.get_ta_state(j, k, the_class=1, polarity=1)))
+        print(" ∧ ".join(l))
+
+    print("\nClause Co-Occurence Matrix:\n")
+    print(tm.clause_co_occurrence(X_test, percentage=True).toarray())
+
+    print("\nLiteral Frequency:\n")
+    print(tm.literal_clause_frequency())
+    experiment_results["literal_frequency"] = tm.literal_clause_frequency().tolist()
+
+    return experiment_results
+
+
+def default_args(**kwargs):
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--epochs", default=250, type=int)
+    parser.add_argument("--number-of-clauses", default=10, type=int)
+    parser.add_argument("--platform", default='CPU', type=str)
+    parser.add_argument("--T", default=10, type=int)
+    parser.add_argument("--s", default=3.0, type=float)
+    parser.add_argument("--number-of-features", default=20, type=int)
+    parser.add_argument("--noise", default=0.1, type=float, help="Noisy XOR")
+    args = parser.parse_args()
+    for key, value in kwargs.items():
+        if key in args.__dict__:
+            setattr(args, key, value)
+    return args
+
+
+if __name__ == "__main__":
+    results = main(default_args())
+    _LOGGER.info(results)
diff --git a/examples/classification/MNISTConvolutionDemo.py b/examples/classification/MNISTConvolutionDemo.py
index f709fc3d..263d65c1 100644
--- a/examples/classification/MNISTConvolutionDemo.py
+++ b/examples/classification/MNISTConvolutionDemo.py
@@ -1,23 +1,21 @@
 import logging
 import argparse
-from data import MNIST
+from tmu.data import MNIST
 from tmu.models.classification.vanilla_classifier import TMClassifier
 from tmu.tools import BenchmarkTimer
 
 _LOGGER = logging.getLogger(__name__)
 
-if __name__ == "__main__":
+def metrics(args):
+    return dict(
+        accuracy=[],
+        train_time=[],
+        test_time=[],
+        args=vars(args)
+    )
 
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--num_clauses", default=8000, type=int)
-    parser.add_argument("--T", default=10000, type=int)
-    parser.add_argument("--s", default=5.0, type=float)
-    parser.add_argument("--max_included_literals", default=32, type=int)
-    parser.add_argument("--device", default="GPU", type=str)
-    parser.add_argument("--weighted_clauses", default=True, type=bool)
-    parser.add_argument("--epochs", default=60, type=int)
-    parser.add_argument("--patch_dim", nargs="+", default=(10, 10), type=int)
-    args = parser.parse_args()
+def main(args):
+    experiment_results = metrics(args)
 
     data = MNIST().get()
 
@@ -31,7 +29,7 @@
         args.s,
         patch_dim=tuple(args.patch_dim),
         max_included_literals=args.max_included_literals,
-        platform=args.device,
+        platform=args.platform,
         weighted_clauses=args.weighted_clauses,
     )
 
@@ -40,11 +38,37 @@
         benchmark1 = BenchmarkTimer(logger=_LOGGER, text="Training Time")
         with benchmark1:
             tm.fit(data["x_train"], data["y_train"])
+        experiment_results["train_time"].append(benchmark1.elapsed())
 
         benchmark2 = BenchmarkTimer(logger=_LOGGER, text="Testing Time")
         with benchmark2:
             result = 100 * (tm.predict(data["x_test"]) == data["y_test"]).mean()
+        experiment_results["train_time"].append(benchmark2.elapsed())
+        experiment_results["accuracy"].append(result)
 
         _LOGGER.info(f"Epoch: {epoch + 1}, Accuracy: {result:.2f}, Training Time: {benchmark1.elapsed():.2f}s, "
                      f"Testing Time: {benchmark2.elapsed():.2f}s")
 
+
+def default_args(**kwargs):
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--num_clauses", default=8000, type=int)
+    parser.add_argument("--T", default=10000, type=int)
+    parser.add_argument("--s", default=5.0, type=float)
+    parser.add_argument("--max_included_literals", default=32, type=int)
+    parser.add_argument("--platform", default="GPU", type=str)
+    parser.add_argument("--weighted_clauses", default=True, type=bool)
+    parser.add_argument("--epochs", default=60, type=int)
+    parser.add_argument("--patch_dim", nargs="+", default=(10, 10), type=int)
+    args = parser.parse_args()
+
+    for key, value in kwargs.items():
+        if key in args.__dict__:
+            setattr(args, key, value)
+
+    return args
+
+
+if __name__ == "__main__":
+    results = main(default_args())
+    _LOGGER.info(results)
diff --git a/examples/classification/MNISTDemo.py b/examples/classification/MNISTDemo.py
index 58fcb999..16efe0a3 100644
--- a/examples/classification/MNISTDemo.py
+++ b/examples/classification/MNISTDemo.py
@@ -1,5 +1,8 @@
 import logging
 import argparse
+
+import numpy as np
+
 from tmu.data import MNIST
 from tmu.models.classification.vanilla_classifier import TMClassifier
 from tmu.tools import BenchmarkTimer
@@ -7,17 +10,16 @@
 
 _LOGGER = logging.getLogger(__name__)
 
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--num_clauses", default=2000, type=int)
-    parser.add_argument("--T", default=5000, type=int)
-    parser.add_argument("--s", default=10.0, type=float)
-    parser.add_argument("--max_included_literals", default=32, type=int)
-    parser.add_argument("--device", default="CPU", type=str, choices=["CPU", "CPU_sparse", "CUDA"])
-    parser.add_argument("--weighted_clauses", default=True, type=bool)
-    parser.add_argument("--epochs", default=60, type=int)
-    args = parser.parse_args()
+def metrics(args):
+    return dict(
+        accuracy=[],
+        train_time=[],
+        test_time=[],
+        args=vars(args)
+    )
 
+def main(args):
+    experiment_results = metrics(args)
     data = MNIST().get()
 
     tm = TMClassifier(
@@ -26,7 +28,7 @@
         T=args.T,
         s=args.s,
         max_included_literals=args.max_included_literals,
-        platform=args.device,
+        platform=args.platform,
         weighted_clauses=args.weighted_clauses,
         seed=42,
     )
@@ -38,17 +40,45 @@
             benchmark1 = BenchmarkTimer(logger=None, text="Training Time")
             with benchmark1:
                 res = tm.fit(
-                    data["x_train"],
-                    data["y_train"],
+                    data["x_train"].astype(np.uint32),
+                    data["y_train"].astype(np.uint32),
                     metrics=["update_p"],
                 )
-            print(res)
+
+            experiment_results["train_time"].append(benchmark1.elapsed())
+
+            # print(res)
             benchmark2 = BenchmarkTimer(logger=None, text="Testing Time")
             with benchmark2:
                 result = 100 * (tm.predict(data["x_test"]) == data["y_test"]).mean()
+                experiment_results["accuracy"].append(result)
+            experiment_results["test_time"].append(benchmark2.elapsed())
 
             _LOGGER.info(f"Epoch: {epoch + 1}, Accuracy: {result:.2f}, Training Time: {benchmark1.elapsed():.2f}s, "
                          f"Testing Time: {benchmark2.elapsed():.2f}s")
 
-        if args.device == "CUDA":
+        if args.platform == "CUDA":
             CudaProfiler().print_timings(benchmark=benchmark_total)
+
+    return experiment_results
+
+
+def default_args(**kwargs):
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--num_clauses", default=2000, type=int)
+    parser.add_argument("--T", default=5000, type=int)
+    parser.add_argument("--s", default=10.0, type=float)
+    parser.add_argument("--max_included_literals", default=32, type=int)
+    parser.add_argument("--platform", default="CPU", type=str, choices=["CPU", "CPU_sparse", "CUDA"])
+    parser.add_argument("--weighted_clauses", default=True, type=bool)
+    parser.add_argument("--epochs", default=60, type=int)
+    args = parser.parse_args()
+    for key, value in kwargs.items():
+        if key in args.__dict__:
+            setattr(args, key, value)
+    return args
+
+
+if __name__ == "__main__":
+    results = main(default_args())
+    _LOGGER.info(results)
diff --git a/examples/classification/MNISTDemoCoalesced.py b/examples/classification/MNISTDemoCoalesced.py
index 4022d571..31530666 100644
--- a/examples/classification/MNISTDemoCoalesced.py
+++ b/examples/classification/MNISTDemoCoalesced.py
@@ -6,17 +6,18 @@
 
 _LOGGER = logging.getLogger(__name__)
 
-if __name__ == "__main__":
+def metrics(args):
+    return dict(
+        number_of_positive_clauses=[],
+        accuracy=[],
+        number_of_includes=[],
+        train_time=[],
+        test_time=[],
+        args=vars(args)
+    )
 
-    default_clauses = 20000 // 5
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--num-clauses", default=default_clauses, type=int)
-    parser.add_argument("--T", default=default_clauses // 4, type=int)
-    parser.add_argument("--s", default=10.0, type=float)
-    parser.add_argument("--weighted-clauses", default=True, type=bool)
-    parser.add_argument("--focused-negative-sampling", default=True, type=bool)
-    parser.add_argument("--epochs", default=60, type=int)
-    args = parser.parse_args()
+def main(args):
+    experiment_results = metrics(args)
 
     data = MNIST().get()
 
@@ -24,6 +25,7 @@
         number_of_clauses=args.num_clauses,
         T=args.T,
         s=args.s,
+        platform=args.platform,
         weighted_clauses=args.weighted_clauses,
         focused_negative_sampling=args.focused_negative_sampling
     )
@@ -34,20 +36,25 @@
         benchmark1 = BenchmarkTimer()
         with benchmark1:
             tm.fit(data["x_train"], data["y_train"])
+        experiment_results["train_time"].append(benchmark1.elapsed())
 
         benchmark2 = BenchmarkTimer()
         with benchmark2:
             result = 100 * (tm.predict(data["x_test"]) == data["y_test"]).mean()
+            experiment_results["accuracy"].append(result)
+        experiment_results["test_time"].append(benchmark2.elapsed())
 
         number_of_positive_clauses = 0
         for i in range(tm.number_of_classes):
             number_of_positive_clauses += (tm.weight_banks[i].get_weights() > 0).sum()
         number_of_positive_clauses /= tm.number_of_classes
+        experiment_results["number_of_positive_clauses"].append(number_of_positive_clauses)
 
         number_of_includes = 0
         for j in range(args.num_clauses):
             number_of_includes += tm.number_of_include_actions(j)
         number_of_includes /= 2 * args.num_clauses
+        experiment_results["number_of_includes"].append(number_of_includes)
 
         _LOGGER.info(
             f"Epoch: {epoch + 1}, "
@@ -57,3 +64,27 @@
             f"Training Time: {benchmark1.elapsed():.2f}s, "
             f"Testing Time: {benchmark2.elapsed():.2f}s"
         )
+
+    return experiment_results
+
+
+def default_args(**kwargs):
+    default_clauses = 20000 // 5
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--num-clauses", default=default_clauses, type=int)
+    parser.add_argument("--T", default=default_clauses // 4, type=int)
+    parser.add_argument("--s", default=10.0, type=float)
+    parser.add_argument("--weighted-clauses", default=True, type=bool)
+    parser.add_argument("--platform", default='CPU', type=str)
+    parser.add_argument("--focused-negative-sampling", default=True, type=bool)
+    parser.add_argument("--epochs", default=60, type=int)
+    args = parser.parse_args()
+    for key, value in kwargs.items():
+        if key in args.__dict__:
+            setattr(args, key, value)
+    return args
+
+
+if __name__ == "__main__":
+    results = main(default_args())
+    _LOGGER.info(results)
diff --git a/examples/classification/MNISTDemoWithSerialization.py b/examples/classification/MNISTDemoWithSerialization.py
index 0e2cc7bd..d260bc5b 100644
--- a/examples/classification/MNISTDemoWithSerialization.py
+++ b/examples/classification/MNISTDemoWithSerialization.py
@@ -5,20 +5,20 @@
 from tmu.tools import BenchmarkTimer
 import pickle
 
-
 _LOGGER = logging.getLogger(__name__)
 
-if __name__ == "__main__":
+def metrics(args):
+    return dict(
+        accuracy=[],
+        train_time=[],
+        test_time=[],
+        serialize_load_time=[],
+        serialize_dump_time=[],
+        args=vars(args)
+    )
 
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--num_clauses", default=50, type=int)
-    parser.add_argument("--T", default=10, type=int)
-    parser.add_argument("--s", default=10.0, type=float)
-    parser.add_argument("--max_included_literals", default=32, type=int)
-    parser.add_argument("--device", default="CPU", type=str)
-    parser.add_argument("--weighted_clauses", default=True, type=bool)
-    parser.add_argument("--epochs", default=60, type=int)
-    args = parser.parse_args()
+def main(args):
+    experiment_results = metrics(args)
 
     data = MNIST().get()
 
@@ -27,7 +27,7 @@
         T=args.T,
         s=args.s,
         max_included_literals=args.max_included_literals,
-        platform=args.device,
+        platform=args.platform,
         weighted_clauses=args.weighted_clauses
     )
 
@@ -36,18 +36,46 @@
         benchmark1 = BenchmarkTimer(logger=_LOGGER, text="Training Time")
         with benchmark1:
             tm.fit(data["x_train"], data["y_train"])
+        experiment_results["train_time"].append(benchmark1.elapsed())
 
         benchmark2 = BenchmarkTimer(logger=_LOGGER, text="Saving model to binary")
         with benchmark2:
             dumped = pickle.dumps(tm)
+        experiment_results["serialize_dump_time"].append(benchmark2.elapsed())
 
         benchmark3 = BenchmarkTimer(logger=_LOGGER, text="Loading model from binary")
         with benchmark3:
             tm2 = pickle.loads(dumped)
+        experiment_results["serialize_load_time"].append(benchmark3.elapsed())
 
         benchmark4 = BenchmarkTimer(logger=_LOGGER, text="Testing Time")
-        with benchmark2:
+        with benchmark4:
             result = 100 * (tm2.predict(data["x_test"]) == data["y_test"]).mean()
+            experiment_results["accuracy"].append(result)
+        experiment_results["test_time"].append(benchmark4.elapsed())
 
         _LOGGER.info(f"Epoch: {epoch + 1}, Accuracy: {result:.2f}, Training Time: {benchmark1.elapsed():.2f}s, "
                      f"Testing Time: {benchmark2.elapsed():.2f}s")
+
+    return experiment_results
+
+
+def default_args(**kwargs):
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--num_clauses", default=50, type=int)
+    parser.add_argument("--T", default=10, type=int)
+    parser.add_argument("--s", default=10.0, type=float)
+    parser.add_argument("--max_included_literals", default=32, type=int)
+    parser.add_argument("--platform", default="CPU", type=str)
+    parser.add_argument("--weighted_clauses", default=True, type=bool)
+    parser.add_argument("--epochs", default=10, type=int)
+    args = parser.parse_args()
+    for key, value in kwargs.items():
+        if key in args.__dict__:
+            setattr(args, key, value)
+    return args
+
+
+if __name__ == "__main__":
+    results = main(default_args())
+    _LOGGER.info(results)
diff --git a/examples/classification/MNISTSparseAbsorbingDemo.py b/examples/classification/MNISTSparseAbsorbingDemo.py
index c564074f..9ce23ee0 100644
--- a/examples/classification/MNISTSparseAbsorbingDemo.py
+++ b/examples/classification/MNISTSparseAbsorbingDemo.py
@@ -1,5 +1,4 @@
 import itertools
-import numpy as np
 import argparse
 import logging
 from tmu.data import MNIST
@@ -8,10 +7,18 @@
 
 _LOGGER = logging.getLogger(__name__)
 
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--epochs", default=250, type=int)
-    args = parser.parse_args()
+def metrics(args):
+    return dict(
+        accuracy=[],
+        accuracy_train=[],
+        absorbed=[],
+        train_time=[],
+        test_time=[],
+        args=vars(args)
+    )
+
+def main(args):
+    experiment_results = metrics(args)
 
     data = MNIST().get()
     X_train = data["x_train"]
@@ -19,24 +26,34 @@
     X_test = data["x_test"]
     y_test = data["y_test"]
 
-    X_train = np.where(X_train.reshape((X_train.shape[0], 28 * 28)) > 75, 1, 0)
-    X_test = np.where(X_test.reshape((X_test.shape[0], 28 * 28)) > 75, 1, 0)
-
-    tm = TMClassifier(2000, 5000, 10.0, platform='CPU_sparse', weighted_clauses=True, literal_sampling=0.1,
-                      max_included_literals=32, feedback_rate_excluded_literals=1, absorbing=100,
-                      literal_insertion_state=-1)
+    tm = TMClassifier(
+        args.num_clauses,
+        args.T,
+        args.s,
+        platform=args.platform,
+        weighted_clauses=args.weighted_clauses,
+        literal_sampling=args.literal_sampling,
+        max_included_literals=args.max_included_literals,
+        feedback_rate_excluded_literals=args.feedback_rate_excluded_literals,
+        absorbing=args.absorbing,
+        literal_insertion_state=args.literal_insertion_state
+    )
 
     _LOGGER.info(f"Running {TMClassifier} for {args.epochs} epochs")
     for e in range(args.epochs):
         benchmark1 = BenchmarkTimer(logger=_LOGGER, text="Training Time")
         with benchmark1:
             tm.fit(X_train, y_train)
+        experiment_results["train_time"].append(benchmark1.elapsed())
 
         benchmark2 = BenchmarkTimer(logger=_LOGGER, text="Testing Time")
         with benchmark2:
             result_test = 100 * (tm.predict(X_test) == y_test).mean()
+            experiment_results["accuracy"].append(result_test)
+        experiment_results["test_time"].append(benchmark2.elapsed())
 
         result_train = 100 * (tm.predict(X_train) == y_train).mean()
+        experiment_results["accuracy_train"].append(result_train)
 
         absorbed = 0.0
         unallocated = 0
@@ -48,7 +65,36 @@
             unallocated += tm.number_of_unallocated_literals(i, j)
 
         absorbed = 100 * absorbed / (num_classes * num_clauses)
+        experiment_results["absorbed"].append(absorbed)
 
         _LOGGER.info(f"Epoch: {e + 1}, Test Accuracy: {result_test:.2f}%, Train Accuracy: {result_train:.2f}%, "
                      f"Absorbed: {absorbed:.2f}%, Unallocated: {unallocated}, Training Time: {benchmark1.elapsed():.2f}s, "
                      f"Testing Time: {benchmark2.elapsed():.2f}s")
+
+        return experiment_results
+
+
+def default_args(**kwargs):
+    parser = argparse.ArgumentParser(description="Configure and run the TMClassifier.")
+    parser.add_argument("--epochs", default=250, type=int)
+    parser.add_argument("--num_clauses", default=2000, type=int)
+    parser.add_argument("--T", default=5000, type=int)
+    parser.add_argument("--s", default=10.0, type=float)
+    parser.add_argument("--platform", default='CPU_sparse', type=str)
+    parser.add_argument("--weighted_clauses", default=True, type=bool)
+    parser.add_argument("--literal_sampling", default=0.1, type=float)
+    parser.add_argument("--max_included_literals", default=32, type=int)
+    parser.add_argument("--feedback_rate_excluded_literals", default=1, type=int)
+    parser.add_argument("--absorbing", default=100, type=int)
+    parser.add_argument("--literal_insertion_state", default=-1, type=int)
+
+    args = parser.parse_args()
+    for key, value in kwargs.items():
+        if key in args.__dict__:
+            setattr(args, key, value)
+    return args
+
+
+if __name__ == "__main__":
+    result = main(default_args())
+    _LOGGER.info(result)
diff --git a/examples/classification/XORDemo.py b/examples/classification/XORDemo.py
index 7d61e48e..3eb34e5e 100644
--- a/examples/classification/XORDemo.py
+++ b/examples/classification/XORDemo.py
@@ -1,22 +1,21 @@
 import logging
 import argparse
-from tmu.data import MNIST, TMUDatasetSource
+from tmu.data.tmu_datasource import TMUDatasetSource
 from tmu.models.classification.vanilla_classifier import TMClassifier
 from tmu.tools import BenchmarkTimer
 
 _LOGGER = logging.getLogger(__name__)
 
-if __name__ == "__main__":
+def metrics(args):
+    return dict(
+        accuracy=[],
+        train_time=[],
+        test_time=[],
+        args=vars(args)
+    )
 
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--num_clauses", default=4, type=int)
-    parser.add_argument("--T", default=10, type=int)
-    parser.add_argument("--s", default=10.0, type=float)
-    parser.add_argument("--max_included_literals", default=32, type=int)
-    parser.add_argument("--device", default="CPU", type=str)
-    parser.add_argument("--weighted_clauses", default=True, type=bool)
-    parser.add_argument("--epochs", default=60, type=int)
-    args = parser.parse_args()
+def main(args):
+    experiment_results = metrics(args)
 
     data = TMUDatasetSource().get_dataset(
         "XOR_biased",
@@ -35,7 +34,7 @@
         T=args.T,
         s=args.s,
         max_included_literals=args.max_included_literals,
-        platform=args.device,
+        platform=args.platform,
         weighted_clauses=args.weighted_clauses
     )
 
@@ -43,11 +42,40 @@
     for epoch in range(args.epochs):
         benchmark1 = BenchmarkTimer(logger=_LOGGER, text="Training Time")
         with benchmark1:
-            tm.fit(data["x_train"], data["y_train"])
+            tm.fit(
+                data["x_train"],
+                data["y_train"]
+            )
+        experiment_results["train_time"].append(benchmark1.elapsed())
 
         benchmark2 = BenchmarkTimer(logger=_LOGGER, text="Testing Time")
         with benchmark2:
             result = 100 * (tm.predict(data["x_test"]) == data["y_test"]).mean()
+            experiment_results["accuracy"].append(result)
+        experiment_results["test_time"].append(benchmark2.elapsed())
 
         _LOGGER.info(f"Epoch: {epoch + 1}, Accuracy: {result:.2f}, Training Time: {benchmark1.elapsed():.2f}s, "
                      f"Testing Time: {benchmark2.elapsed():.2f}s")
+
+    return experiment_results
+
+def default_args(**kwargs):
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--num_clauses", default=4, type=int)
+    parser.add_argument("--T", default=10, type=int)
+    parser.add_argument("--s", default=10.0, type=float)
+    parser.add_argument("--max_included_literals", default=32, type=int)
+    parser.add_argument("--platform", default="CPU", type=str)
+    parser.add_argument("--weighted_clauses", default=True, type=bool)
+    parser.add_argument("--epochs", default=60, type=int)
+    args = parser.parse_args()
+    for key, value in kwargs.items():
+        if key in args.__dict__:
+            setattr(args, key, value)
+    return args
+
+if __name__ == "__main__":
+
+
+    score = main(default_args())
+    _LOGGER.info(score)
diff --git a/examples/classification/__init__.py b/examples/classification/__init__.py
index e69de29b..27171da3 100644
--- a/examples/classification/__init__.py
+++ b/examples/classification/__init__.py
@@ -0,0 +1,11 @@
+import os
+import glob
+
+__all__ = []
+directory = os.path.dirname(__file__)
+
+for filename in glob.glob(os.path.join(directory, '*.py')):
+    module = os.path.basename(filename)[:-3]  # Remove .py extension
+    if module != '__init__':
+        __import__(__name__ + '.' + module)
+        __all__.append(module)
\ No newline at end of file
diff --git a/examples/composite/TMCompositeCIFAR10Demo.py b/examples/composite/TMCompositeCIFAR10Demo.py
new file mode 100644
index 00000000..50e6ac99
--- /dev/null
+++ b/examples/composite/TMCompositeCIFAR10Demo.py
@@ -0,0 +1,151 @@
+import argparse
+import logging
+
+from tmu.composite.callbacks.base import TMCompositeCallback
+from tmu.composite.composite import TMComposite
+from tmu.composite.components.adaptive_thresholding import AdaptiveThresholdingComponent
+from tmu.composite.components.color_thermometer_scoring import ColorThermometerComponent
+from tmu.composite.components.histogram_of_gradients import HistogramOfGradientsComponent
+from tmu.composite.config import TMClassifierConfig
+from tmu.data.cifar10 import CIFAR10
+from tmu.models.classification.vanilla_classifier import TMClassifier
+import pathlib
+
+_LOGGER = logging.getLogger(__name__)
+
+def metrics(args):
+    return dict(
+        accuracy=[]
+    )
+
+def main(args):
+    experiment_results = metrics(args)
+
+    platform = args.platform
+    epochs = args.epochs
+    checkpoint_path = pathlib.Path("checkpoints")
+    checkpoint_path.mkdir(exist_ok=True)
+
+    composite_path = checkpoint_path / "composite"
+    composite_path.mkdir(exist_ok=True)
+
+    component_path = checkpoint_path / "components"
+    component_path.mkdir(exist_ok=True)
+
+    data = CIFAR10().get()
+    X_train_org = data["x_train"]
+    Y_train = data["y_train"]
+    X_test_org = data["x_test"]
+    Y_test = data["y_test"]
+
+    data_train = dict(
+        X=X_train_org,
+        Y=Y_train
+    )
+
+    data_test = dict(
+        X=X_test_org,
+        Y=Y_test
+    )
+
+    class TMCompositeCheckpointCallback(TMCompositeCallback):
+
+        def on_epoch_component_begin(self, component, epoch, logs=None):
+            pass
+
+        def on_epoch_component_end(self, component, epoch, logs=None):
+            component.save(component_path / f"{component}-{epoch}.pkl")
+
+    class TMCompositeEvaluationCallback(TMCompositeCallback):
+
+        def __init__(self, data):
+            super().__init__()
+            self.best_acc = 0.0
+            self.data = data
+
+    # def on_epoch_end(self, composite, epoch, logs=None):
+    #     preds = composite.predict(data=self.data)
+    #     print("Team Accuracy: %.1f" % (100 * (preds == self.data["Y"]).mean()))
+
+    # Define the composite model
+    composite_model = TMComposite(
+        components=[
+            AdaptiveThresholdingComponent(TMClassifier, TMClassifierConfig(
+                number_of_clauses=2000,
+                T=500,
+                s=10.0,
+                max_included_literals=32,
+                platform=platform,
+                weighted_clauses=True,
+                patch_dim=(10, 10),
+            ), epochs=epochs),
+
+            ColorThermometerComponent(TMClassifier, TMClassifierConfig(
+                number_of_clauses=2000,
+                T=1500,
+                s=2.5,
+                max_included_literals=32,
+                platform=platform,
+                weighted_clauses=True,
+                patch_dim=(3, 3),
+            ), resolution=8, epochs=epochs),
+
+            ColorThermometerComponent(TMClassifier, TMClassifierConfig(
+                number_of_clauses=2000,
+                T=1500,
+                s=2.5,
+                max_included_literals=32,
+                platform=platform,
+                weighted_clauses=True,
+                patch_dim=(4, 4),
+            ), resolution=8, epochs=epochs),
+
+            HistogramOfGradientsComponent(TMClassifier, TMClassifierConfig(
+                number_of_clauses=2000,
+                T=50,
+                s=10.0,
+                max_included_literals=32,
+                platform=platform,
+                weighted_clauses=False
+            ), epochs=epochs)
+        ],
+        use_multiprocessing=False
+    )
+
+    # Train the composite model
+    composite_model.fit(
+        data=data_train,
+        callbacks=[
+            TMCompositeCheckpointCallback(),
+            TMCompositeEvaluationCallback(data=data_test)
+        ]
+    )
+
+    preds = composite_model.predict(data=data_test)
+
+    y_true = data_test["Y"].flatten()
+    for k, v in preds.items():
+        comp_acc = 100 * (v == y_true).mean()
+        print(f"{k} Accuracy: %.1f" % (comp_acc,))
+
+        if k not in experiment_results["accuracy"]:
+            experiment_results["accuracy"][k] = []
+        experiment_results["accuracy"][k].append(comp_acc)
+
+    return experiment_results
+
+
+def default_args(**kwargs):
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--platform", default="CPU", type=str)
+    parser.add_argument("--epochs", default=30, type=int)
+    args = parser.parse_args()
+    for key, value in kwargs.items():
+        if key in args.__dict__:
+            setattr(args, key, value)
+    return args
+
+
+if __name__ == "__main__":
+    results = main(default_args())
+    _LOGGER.info(results)
diff --git a/examples/composite/__init__.py b/examples/composite/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/examples/composite/hpsearch/__init__.py b/examples/composite/hpsearch/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/examples/composite/hpsearch/example_2.py b/examples/composite/hpsearch/example_2.py
new file mode 100644
index 00000000..c086e7dd
--- /dev/null
+++ b/examples/composite/hpsearch/example_2.py
@@ -0,0 +1,69 @@
+from tmu.composite.tuner import TMCompositeTuner
+from tmu.data.cifar10 import CIFAR10
+import numpy as np
+
+
+def create_subset(X, Y, samples_per_class):
+    """
+    Create a balanced subset of the data with an equal number of samples per class.
+
+    Parameters:
+    - X: numpy array, feature matrix with shape (n_samples, n_features)
+    - Y: numpy array, labels with shape (n_samples,) where each element is an integer class label
+    - samples_per_class: int, the number of samples to include from each class
+
+    Returns:
+    - A tuple of two numpy arrays: (subset_X, subset_Y) shuffled and balanced subset.
+    """
+    unq = len(np.unique(Y))
+
+    subset_X, subset_Y = [], []
+
+    for i in range(unq):  # Assuming 10 classes in CIFAR-10
+        mask = (Y == i).reshape(-1)
+        subset_X.append(X[mask][:samples_per_class])
+        subset_Y.append(Y[mask][:samples_per_class])
+
+    # Using np.concatenate for clarity
+    subset_X, subset_Y = np.concatenate(subset_X), np.concatenate(subset_Y)
+
+    # Shuffling the dataset
+    indices = np.arange(subset_X.shape[0])
+    np.random.shuffle(indices)
+
+    return subset_X[indices], subset_Y[indices]
+
+
+if __name__ == "__main__":
+    data = CIFAR10().get()
+    X_train_org = data["x_train"]
+    Y_train = data["y_train"]
+    X_test_org = data["x_test"]
+    Y_test = data["y_test"]
+    percentage = 0.1
+    num_train_samples = int(5000 * percentage)
+    num_test_samples = int(1000 * percentage)
+
+    X_train_subset, Y_train_subset = create_subset(X_train_org, Y_train, num_train_samples)
+    X_test_subset, Y_test_subset = create_subset(X_test_org, Y_test, num_test_samples)
+    Y_test_subset = Y_test_subset.reshape(Y_test_subset.shape[0])
+    Y_train_subset = Y_train_subset.reshape(Y_train_subset.shape[0])
+    data_train = dict(X=X_train_subset, Y=Y_train_subset)
+    data_test = dict(X=X_test_subset, Y=Y_test_subset)
+
+    # Instantiate tuner
+    tuner = TMCompositeTuner(
+        data_train=data_train,
+        data_test=data_test,
+        n_jobs=1  # for parallelization; set to 1 for no parallelization
+    )
+
+    # Specify number of trials (iterations of the tuning process)
+    n_trials = 100
+
+    # Run the tuner
+    best_params, best_value = tuner.tune(n_trials=n_trials)
+
+    # Print out the results
+    print("Best Parameters:", best_params)
+    print("Best Value:", best_value)
diff --git a/examples/experimental/classification/InterpretabilityDemoAND.py b/examples/experimental/classification/InterpretabilityDemoAND.py
index cdb524d3..7b57b1fe 100644
--- a/examples/experimental/classification/InterpretabilityDemoAND.py
+++ b/examples/experimental/classification/InterpretabilityDemoAND.py
@@ -1,4 +1,4 @@
-from tmu.models.classification.multichannel_classifier import TMMultiChannelClassifier
+from tmu.experimental.models.multichannel_classifier import TMMultiChannelClassifier
 import numpy as np
 
 number_of_and_components = 2
diff --git a/examples/experimental/classification/InterpretabilityDemoOneVsOne.py b/examples/experimental/classification/InterpretabilityDemoOneVsOne.py
index 16a0fe8f..cfddaae7 100644
--- a/examples/experimental/classification/InterpretabilityDemoOneVsOne.py
+++ b/examples/experimental/classification/InterpretabilityDemoOneVsOne.py
@@ -1,6 +1,5 @@
-from tmu.models.classification.one_vs_one_classifier import TMOneVsOneClassifier
+from tmu.experimental.models.one_vs_one_classifier import TMOneVsOneClassifier
 import numpy as np
-from time import time
 
 clauses = 4
 number_of_features = 20
diff --git a/examples/experimental/classification/MNISTDemo2DConvolutionOneVsOne.py b/examples/experimental/classification/MNISTDemo2DConvolutionOneVsOne.py
index 14251a47..3d063461 100644
--- a/examples/experimental/classification/MNISTDemo2DConvolutionOneVsOne.py
+++ b/examples/experimental/classification/MNISTDemo2DConvolutionOneVsOne.py
@@ -3,7 +3,7 @@
 
 from keras.datasets import mnist
 
-from tmu.models.classification.one_vs_one_classifier import TMOneVsOneClassifier
+from tmu.experimental.models.one_vs_one_classifier import TMOneVsOneClassifier
 
 (X_train, Y_train), (X_test, Y_test) = mnist.load_data()
 
diff --git a/examples/experimental/classification/MNISTDemoOneVsOne.py b/examples/experimental/classification/MNISTDemoOneVsOne.py
index dd968c2f..116b702a 100644
--- a/examples/experimental/classification/MNISTDemoOneVsOne.py
+++ b/examples/experimental/classification/MNISTDemoOneVsOne.py
@@ -3,7 +3,7 @@
 
 from keras.datasets import mnist
 
-from tmu.models.classification.one_vs_one_classifier import TMOneVsOneClassifier
+from tmu.experimental.models.one_vs_one_classifier import TMOneVsOneClassifier
 
 (X_train, Y_train), (X_test, Y_test) = mnist.load_data()
 
diff --git a/examples/experimental/relational/RelationalTMDemo.py b/examples/experimental/relational/RelationalTMDemo.py
index 1aad2701..668c5f78 100644
--- a/examples/experimental/relational/RelationalTMDemo.py
+++ b/examples/experimental/relational/RelationalTMDemo.py
@@ -1,8 +1,4 @@
-import numpy as np
-import keras
-from time import time
-
-from tmu.models.relational.vanilla_relational import TMRelational
+from tmu.experimental.models.relational.vanilla_relational import TMRelational
 
 X = [
 		((("Parent", "Mary", "Bob"), True), (("Parent", "Bob", "Peter"), True), (("Ancestor", "Mary", "Bob"), True), (("Ancestor", "Bob", "Peter"), True), (("Ancestor", "Mary", "Peter"), True)),
diff --git a/examples/type_iii_feedback/coma.ipynb b/examples/experimental/type_iii_feedback/coma.ipynb
similarity index 100%
rename from examples/type_iii_feedback/coma.ipynb
rename to examples/experimental/type_iii_feedback/coma.ipynb
diff --git a/examples/type_iii_feedback/exp2.ipynb b/examples/experimental/type_iii_feedback/exp2.ipynb
similarity index 100%
rename from examples/type_iii_feedback/exp2.ipynb
rename to examples/experimental/type_iii_feedback/exp2.ipynb
diff --git a/examples/experimental/type_iii_feedback/utils.py b/examples/experimental/type_iii_feedback/utils.py
new file mode 100644
index 00000000..e69de29b
diff --git a/examples/regression/RegressionDemo.py b/examples/regression/RegressionDemo.py
index 430b1a3f..97591092 100644
--- a/examples/regression/RegressionDemo.py
+++ b/examples/regression/RegressionDemo.py
@@ -9,44 +9,69 @@
 
 _LOGGER = logging.getLogger(__name__)
 
-if __name__ == "__main__":
+def metrics(args):
+    return dict(
+        rmsd=[],
+        train_time=[],
+        test_time=[]
+    )
+
+def main(args):
+    experiment_results = metrics(args)
+
+    california_housing = datasets.fetch_california_housing()
+    X = california_housing.data
+    Y = california_housing.target
+
+    b = StandardBinarizer(max_bits_per_feature=10)
+    X_transformed = b.fit_transform(X)
+
+    tm = TMRegressor(
+        args.num_clauses,
+        args.T,
+        args.s,
+        platform=args.platform,
+        weighted_clauses=args.weighted_clauses
+    )
+
+    tm_results = np.empty(0)
+
+    _LOGGER.info(f"Running RMSD with {TMRegressor} for {args.num_runs}")
+    for run in range(args.num_runs):
+        X_train, X_test, Y_train, Y_test = train_test_split(X_transformed, Y)
 
-	parser = argparse.ArgumentParser()
-	parser.add_argument("--num_clauses", default=1000, type=int)
-	parser.add_argument("--T", default=500 * 10, type=int)
-	parser.add_argument("--s", default=2.75, type=float)
-	parser.add_argument("--device", default="CPU", type=str)
-	parser.add_argument("--weighted_clauses", default=True, type=bool)
-	parser.add_argument("--epochs", default=30, type=int)
-	parser.add_argument("--num-runs", default=25, type=int)
-	args = parser.parse_args()
-
-	california_housing = datasets.fetch_california_housing()
-	X = california_housing.data
-	Y = california_housing.target
-
-	b = StandardBinarizer(max_bits_per_feature=10)
-	X_transformed = b.fit_transform(X)
-
-	tm = TMRegressor(
-		args.num_clauses,
-		args.T,
-		args.s,
-		platform=args.device,
-		weighted_clauses=args.weighted_clauses
-	)
-
-	tm_results = np.empty(0)
-
-	_LOGGER.info(f"Running RMSD with {TMRegressor} for {args.num_runs}")
-	for run in range(args.num_runs):
-		X_train, X_test, Y_train, Y_test = train_test_split(X_transformed, Y)
-
-		benchmark1 = BenchmarkTimer(logger=_LOGGER, text="Training Time")
-		with benchmark1:
-			for epoch in range(args.epochs):
-				tm.fit(X_train, Y_train)
-
-		tm_results = np.append(tm_results, np.sqrt(((tm.predict(X_test) - Y_test) ** 2).mean()))
-
-		_LOGGER.info(f"#{run + 1}: RMSD: {tm_results.mean()} +/- {1.96 * tm_results.std() / np.sqrt(run + 1)} ({benchmark1.elapsed()})")
+        for epoch in range(args.epochs):
+            benchmark1 = BenchmarkTimer(logger=_LOGGER, text="Training Time")
+            with benchmark1:
+                tm.fit(X_train, Y_train)
+            experiment_results["train_time"].append(benchmark1.elapsed())
+
+            benchmark2 = BenchmarkTimer(logger=_LOGGER, text="Testing Time")
+            with benchmark2:
+                tm_results = np.append(tm_results, np.sqrt(((tm.predict(X_test) - Y_test) ** 2).mean()))
+            experiment_results["test_time"].append(benchmark1.elapsed())
+            experiment_results["rmsd"].append(tm_results.mean())
+
+            _LOGGER.info(
+                f"#{run + 1} - Epoch: {epoch}: RMSD: {tm_results.mean()} +/- {1.96 * tm_results.std() / np.sqrt(run + 1)} ({benchmark1.elapsed()})")
+
+
+def default_args(**kwargs):
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--num_clauses", default=1000, type=int)
+    parser.add_argument("--T", default=500 * 10, type=int)
+    parser.add_argument("--s", default=2.75, type=float)
+    parser.add_argument("--platform", default="CPU", type=str)
+    parser.add_argument("--weighted_clauses", default=True, type=bool)
+    parser.add_argument("--epochs", default=30, type=int)
+    parser.add_argument("--num-runs", default=25, type=int)
+    args = parser.parse_args()
+    for key, value in kwargs.items():
+        if key in args.__dict__:
+            setattr(args, key, value)
+    return args
+
+
+if __name__ == "__main__":
+    results = main(default_args())
+    _LOGGER.info(results)
diff --git a/examples/relational/RelationalTMDemo.py b/examples/relational/RelationalTMDemo.py
deleted file mode 100644
index 7a6f0ce9..00000000
--- a/examples/relational/RelationalTMDemo.py
+++ /dev/null
@@ -1,59 +0,0 @@
-import argparse
-
-from tmu.models.relational.vanilla_relational import TMRelational
-import logging
-
-_LOGGER = logging.getLogger(__name__)
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--num_clauses", default=20, type=int)
-    parser.add_argument("--T", default=40, type=int)
-    parser.add_argument("--s", default=5.0, type=float)
-    parser.add_argument("--max-included-literals", default=3, type=int)
-    parser.add_argument("--feature-negation", default=False, type=bool)
-    parser.add_argument("--output-balancing", default=True, type=bool)
-    parser.add_argument("--device", default="CPU", type=str)
-    parser.add_argument("--epochs", default=60, type=int)
-    args = parser.parse_args()
-
-    X = [
-        (
-            (("Parent", "Mary", "Bob"), True),
-            (("Parent", "Bob", "Peter"), True),
-            (("Ancestor", "Mary", "Bob"), True),
-            (("Ancestor", "Bob", "Peter"), True),
-            (("Ancestor", "Mary", "Peter"), True)),
-        (
-            (("Parent", "Ida", "Chris"), True),
-            (("Parent", "Chris", "Ann"), True),
-            (("Ancestor", "Ida", "Chris"), True),
-            (("Ancestor", "Chris", "Ann"), True),
-            (("Ancestor", "Ida", "Ann"), True)
-        )
-    ]
-
-    output_active = [
-        ("Ancestor", "Ida", "Chris"),
-        ("Ancestor", "Chris", "Ann"),
-        ("Ancestor", "Ida", "Ann"),
-        ("Ancestor", "Mary", "Bob"),
-        ("Ancestor", "Bob", "Peter"),
-        ("Ancestor", "Mary", "Peter")
-    ]
-
-    tm = TMRelational(
-        number_of_clauses=args.num_clauses,
-        T=args.T,
-        s=args.s,
-        output_active_facts=output_active,
-        max_included_literals=args.max_included_literals,
-        feature_negation=args.feature_negation,
-        platform=args.device,
-        output_balancing=args.output_balancing
-    )
-
-    _LOGGER.info(f"Accuracy over {args.epochs} epochs.")
-    for epoch in range(args.epochs):
-        tm.fit(X)
-        _LOGGER.info(f"Epoch: {epoch}")
diff --git a/examples/requirements.txt b/examples/requirements.txt
deleted file mode 100644
index 480638e1..00000000
--- a/examples/requirements.txt
+++ /dev/null
@@ -1,2 +0,0 @@
-scikit-learn
-tensorflow
diff --git a/examples/stats.py b/examples/stats.py
deleted file mode 100644
index 97f36ad8..00000000
--- a/examples/stats.py
+++ /dev/null
@@ -1,27 +0,0 @@
-import numpy as np
-import sys
-
-m = np.loadtxt(sys.argv[1])
-start = int(sys.argv[2])
-
-print("MIN", m[np.logical_and(m[:,1] >= start, m[:,1] < start+25),:].min(axis=0))
-print("MEAN", m[np.logical_and(m[:,1] >= start, m[:,1] < start+25),:].mean(axis=0))
-print("MEDIAN", np.median(m[np.logical_and(m[:,1] >= start, m[:,1] < start+25),:], axis=0))
-print("STD", m[np.logical_and(m[:,1] >= start, m[:,1] < start+25),:].std(axis=0))
-print("95%", np.percentile(m[np.logical_and(m[:,1] >= start, m[:,1] < start+25),:], 95, axis=0))
-print("MAX", m[np.logical_and(m[:,1] >= start, m[:,1] < start+25),:].max(axis=0))
-print()
-
-ids = m[:,0]
-data = m[:,1:]
-_ndx = np.argsort(ids)
-_id, _pos  = np.unique(ids[_ndx], return_index=True)
-g_max = np.maximum.reduceat(data[_ndx], _pos)
-g_min = np.minimum.reduceat(data[_ndx], _pos)
-print("MAX-MIN:", g_min.max(axis=0))
-print("AVERAGE MIN", g_min.mean(axis=0))
-print("VARIANCE MIN", g_min.var(axis=0))
-print()
-print("MIN-MAX:", g_max.min(axis=0))
-print("AVERAGE MAX", g_max.mean(axis=0))
-print("VARIANCE MAX", g_max.var(axis=0))
\ No newline at end of file
diff --git a/pyproject.toml b/pyproject.toml
index 26f58dca..c920d9b8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -2,7 +2,7 @@
 requires = [
     "setuptools",
     "wheel",
-    "cffi>=1.0.0",
+    "cffi>=1.15.0",
     "tomli"
 ]
 build-backend = "setuptools.build_meta"
@@ -18,15 +18,34 @@ authors = [
     { email = "per.andersen@uia.no", name = "Per-Arne Andersen" }
 ]
 dependencies = [
-    "cffi>=1.0.0",
+    "cffi>=1.15.0",
     "numpy",
     "pandas",
     "scikit-learn",
     "tqdm",
-    "requests"
+    "requests",
+    "xxhash"
 ]
 requires-python = ">=3.10"
 
+[project.optional-dependencies]
+composite = [
+    "opencv-python-headless",
+    "pydantic",
+    "optuna",
+    "joblib"
+]
+
+examples = [
+    "scikit-learn",
+    "tensorflow"
+]
+
+tests = [
+    "pytest"
+]
+
+
 # Include additional files
 [project.urls]
 homepage = "https://github.com/cair/tmu/"
diff --git a/scripts/example_verifier/__init__.py b/scripts/example_verifier/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/scripts/example_verifier/verify_experiments.py b/scripts/example_verifier/verify_experiments.py
new file mode 100644
index 00000000..b9fbe0b0
--- /dev/null
+++ b/scripts/example_verifier/verify_experiments.py
@@ -0,0 +1,80 @@
+import json
+import os
+from multiprocessing import Pool
+
+
+import examples.classification
+import examples.autoencoder
+import examples.regression
+import examples.composite
+import inspect
+
+
+def run_module_main_with_args(module_info):
+    module_name, fns = module_info
+    if "args_fn" in fns and "main_fn" in fns and fns["args_fn"] and fns["main_fn"]:
+        print(f"Running {module_name}...")
+        args = fns["args_fn"](epochs=1)
+        result = fns["main_fn"](args)
+        return module_name, result
+    else:
+        print(f"Skipping {module_name}, missing args_fn or main_fn.")
+        return module_name, None
+
+
+def gather_module_info(packages):
+    modules_dict = {}
+    for package in packages:
+        public_attributes = [attr for attr in dir(package) if not attr.startswith('_')]
+
+        for attr in public_attributes:
+            module = getattr(package, attr)
+            if inspect.ismodule(module):
+                module_key = f"{package.__name__}.{attr}"
+                modules_dict[module_key] = {}
+                if hasattr(module, 'default_args') and inspect.isfunction(getattr(module, 'default_args')):
+                    modules_dict[module_key]['args_fn'] = getattr(module, 'default_args')
+                if hasattr(module, 'main') and inspect.isfunction(getattr(module, 'main')):
+                    modules_dict[module_key]['main_fn'] = getattr(module, 'main')
+
+    # Filter out modules that don't contain either default_args or main
+    modules_dict = {k: v for k, v in modules_dict.items() if 'args_fn' in v and 'main_fn' in v}
+    return modules_dict
+
+
+if __name__ == "__main__":
+    import multiprocessing
+    multiprocessing.set_start_method("spawn")
+    packages = [
+        examples.classification,
+        examples.autoencoder,
+        examples.regression,
+        examples.composite
+    ]
+    modules_dict = gather_module_info(packages)
+    modules_info = list(modules_dict.items())
+
+    for cls, fns in modules_dict.items():
+        args = fns["args_fn"](epochs=1)  # Attempt to override epochs
+        assert args.epochs == 1, "We should be able to override args"
+
+    modules_info = list(modules_dict.items())
+
+    results_json = {}
+
+    # Setup multiprocessing pool
+    with Pool(processes=os.cpu_count()) as pool:
+        for module_name, result in pool.imap_unordered(run_module_main_with_args, modules_info):
+            if result is not None:
+                results_json[module_name] = result
+                # Save/update results incrementally
+                with open('module_results.json', 'w') as f:
+                    json.dump(results_json, f, indent=4)
+                print(f"Results saved for {module_name}.")
+
+
+    # Optionally, save the results to a JSON file
+    with open('module_results.json', 'w') as f:
+        json.dump(results_json, f, indent=4)
+
+    print("Results saved to module_results.json.")
diff --git a/tmu/__init__.py b/tmu/__init__.py
index 556e4a70..63a4d4b0 100644
--- a/tmu/__init__.py
+++ b/tmu/__init__.py
@@ -21,4 +21,4 @@
 except ImportError as e:
     raise ImportError("Could not import cffi compiled libraries. To fix this problem, run pip install -e .", e)
 
-__version__ = "0.8.2"
+__version__ = "0.8.3"
diff --git a/tmu/clause_bank/clause_bank_cuda.py b/tmu/clause_bank/clause_bank_cuda.py
index 279f56b0..ac59093c 100644
--- a/tmu/clause_bank/clause_bank_cuda.py
+++ b/tmu/clause_bank/clause_bank_cuda.py
@@ -388,7 +388,7 @@ def produce_autoencoder_example(
         for key, value in kwargs.items():
             if key not in self._logged_unknown_args:
                 self._logged_unknown_args.add(key)
-                _LOGGER.error(
+                _LOGGER.warning(
                     f"Unknown positional argument for {self}: argument_name={key}, argument_value={value}, class={type(self)}")
 
         (
diff --git a/tmu/clause_bank/clause_bank_sparse.py b/tmu/clause_bank/clause_bank_sparse.py
index 37d5cf90..2a9830ec 100644
--- a/tmu/clause_bank/clause_bank_sparse.py
+++ b/tmu/clause_bank/clause_bank_sparse.py
@@ -97,42 +97,42 @@ def __init__(
 
     def _cffi_init(self):
         self.clause_bank_included = np.ascontiguousarray(np.zeros((self.number_of_clauses, self.number_of_literals, 2),
-                                                                  dtype=np.uint16))  # Contains index and state of included literals, none at start
-        self.ptr_clause_bank_included = ffi.cast("unsigned short *", self.clause_bank_included.ctypes.data)
-        self.clause_bank_included_length = np.ascontiguousarray(np.zeros(self.number_of_clauses, dtype=np.uint16))
-        self.ptr_clause_bank_included_length = ffi.cast("unsigned short *",
+                                                                  dtype=np.uint32))  # Contains index and state of included literals, none at start
+        self.ptr_clause_bank_included = ffi.cast("unsigned int *", self.clause_bank_included.ctypes.data)
+        self.clause_bank_included_length = np.ascontiguousarray(np.zeros(self.number_of_clauses, dtype=np.uint32))
+        self.ptr_clause_bank_included_length = ffi.cast("unsigned int *",
                                                         self.clause_bank_included_length.ctypes.data)
 
         self.clause_bank_included_absorbed = np.ascontiguousarray(
             np.zeros((self.number_of_clauses, self.number_of_literals),
-                     dtype=np.uint16))  # Contains index and state of included literals, none at start
-        self.cbia_p = ffi.cast("unsigned short *", self.clause_bank_included_absorbed.ctypes.data)
+                     dtype=np.uint32))  # Contains index and state of included literals, none at start
+        self.cbia_p = ffi.cast("unsigned int *", self.clause_bank_included_absorbed.ctypes.data)
         self.clause_bank_included_absorbed_length = np.ascontiguousarray(
-            np.zeros(self.number_of_clauses, dtype=np.uint16))
-        self.cbial_p = ffi.cast("unsigned short *", self.clause_bank_included_absorbed_length.ctypes.data)
+            np.zeros(self.number_of_clauses, dtype=np.uint32))
+        self.cbial_p = ffi.cast("unsigned int *", self.clause_bank_included_absorbed_length.ctypes.data)
 
         self.clause_bank_excluded = np.ascontiguousarray(np.zeros((self.number_of_clauses, self.number_of_literals, 2),
-                                                                  dtype=np.uint16))  # Contains index and state of excluded literals
-        self.ptr_clause_bank_excluded = ffi.cast("unsigned short *", self.clause_bank_excluded.ctypes.data)
+                                                                  dtype=np.uint32))  # Contains index and state of excluded literals
+        self.ptr_clause_bank_excluded = ffi.cast("unsigned int *", self.clause_bank_excluded.ctypes.data)
         self.clause_bank_excluded_length = np.ascontiguousarray(
-            np.zeros(self.number_of_clauses, dtype=np.uint16))  # All literals excluded at start
-        self.ptr_clause_bank_excluded_length = ffi.cast("unsigned short *",
+            np.zeros(self.number_of_clauses, dtype=np.uint32))  # All literals excluded at start
+        self.ptr_clause_bank_excluded_length = ffi.cast("unsigned int *",
                                                         self.clause_bank_excluded_length.ctypes.data)
         self.clause_bank_excluded_length[:] = int(self.number_of_literals * self.literal_sampling)
 
         for j in range(self.number_of_clauses):
-            literal_indexes = np.arange(self.number_of_literals, dtype=np.uint16)
+            literal_indexes = np.arange(self.number_of_literals, dtype=np.uint32)
             self.rng.shuffle(literal_indexes)
             self.clause_bank_excluded[j, :, 0] = literal_indexes  # Initialize clause literals randomly
             self.clause_bank_excluded[j, :,
             1] = self.number_of_states // 2 - 1  # Initialize excluded literals in least forgotten state
 
         self.clause_bank_unallocated = np.ascontiguousarray(np.zeros((self.number_of_clauses, self.number_of_literals),
-                                                                     dtype=np.uint16))  # Contains index and unallocated literals
-        self.ptr_clause_bank_unallocated = ffi.cast("unsigned short *", self.clause_bank_unallocated.ctypes.data)
+                                                                     dtype=np.uint32))  # Contains index and unallocated literals
+        self.ptr_clause_bank_unallocated = ffi.cast("unsigned int *", self.clause_bank_unallocated.ctypes.data)
         self.clause_bank_unallocated_length = np.ascontiguousarray(
-            np.zeros(self.number_of_clauses, dtype=np.uint16))  # All literals excluded at start
-        self.ptr_clause_bank_unallocated_length = ffi.cast("unsigned short *",
+            np.zeros(self.number_of_clauses, dtype=np.uint32))  # All literals excluded at start
+        self.ptr_clause_bank_unallocated_length = ffi.cast("unsigned int *",
                                                            self.clause_bank_unallocated_length.ctypes.data)
         self.clause_bank_unallocated_length[:] = self.number_of_literals - int(
             self.number_of_literals * self.literal_sampling)
diff --git a/tmu/clause_bank/cuda/clause_feedback.cu b/tmu/clause_bank/cuda/clause_feedback.cu
index bfe5c687..22af021c 100644
--- a/tmu/clause_bank/cuda/clause_feedback.cu
+++ b/tmu/clause_bank/cuda/clause_feedback.cu
@@ -51,9 +51,6 @@ extern "C"
 
         carry = active;
         for (int b = 0; b < number_of_state_bits; ++b) {
-            if (carry == 0)
-                break;
-
             carry_next = ta_state[b] & carry; // Sets carry bits (overflow) passing on to next bit
             ta_state[b] = ta_state[b] ^ carry; // Performs increments with XOR
             carry = carry_next;
@@ -73,9 +70,6 @@ extern "C"
 
         carry = active;
         for (int b = 0; b < number_of_state_bits; ++b) {
-            if (carry == 0)
-                break;
-
             carry_next = (~ta_state[b]) & carry; // Sets carry bits (overflow) passing on to next bit
             ta_state[b] = ta_state[b] ^ carry; // Performs increments with XOR
             carry = carry_next;
diff --git a/tmu/composite/__init__.py b/tmu/composite/__init__.py
new file mode 100644
index 00000000..a68927d6
--- /dev/null
+++ b/tmu/composite/__init__.py
@@ -0,0 +1,1 @@
+__version__ = "0.1.0"
\ No newline at end of file
diff --git a/tmu/composite/callbacks/__init__.py b/tmu/composite/callbacks/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tmu/composite/callbacks/base.py b/tmu/composite/callbacks/base.py
new file mode 100644
index 00000000..33dd107e
--- /dev/null
+++ b/tmu/composite/callbacks/base.py
@@ -0,0 +1,37 @@
+from multiprocessing import Queue
+
+
+class TMCompositeCallback:
+
+    def __init__(self):
+        pass
+
+    def on_epoch_component_begin(self, component, epoch, logs=None):
+        pass
+
+    def on_epoch_component_end(self, component, epoch, logs=None):
+        pass
+
+    def on_train_composite_end(self, composite, logs=None):
+        pass
+
+    def on_train_composite_begin(self, composite, logs=None):
+        pass
+
+
+class TMCompositeCallbackProxy:
+
+    def __init__(self, queue: Queue):
+        self.queue = queue
+
+    def on_epoch_component_begin(self, component, epoch, logs=None):
+        self.queue.put(('on_epoch_component_begin', component, epoch, logs))
+
+    def on_epoch_component_end(self, component, epoch, logs=None):
+        self.queue.put(('on_epoch_component_end', component, epoch, logs))
+
+    def on_train_composite_end(self, composite, logs=None):
+        self.queue.put(('on_train_composite_end', composite, logs))
+
+    def on_train_composite_begin(self, composite, logs=None):
+        self.queue.put(('on_train_composite_begin', composite, logs))
diff --git a/tmu/composite/components/__init__.py b/tmu/composite/components/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tmu/composite/components/adaptive_thresholding.py b/tmu/composite/components/adaptive_thresholding.py
new file mode 100644
index 00000000..87becb67
--- /dev/null
+++ b/tmu/composite/components/adaptive_thresholding.py
@@ -0,0 +1,24 @@
+import numpy as np
+import cv2
+from tmu.composite.components.base import TMComponent
+
+
+class AdaptiveThresholdingComponent(TMComponent):
+
+    def __init__(self, model_cls, model_config, **kwargs) -> None:
+        super().__init__(model_cls=model_cls, model_config=model_config, **kwargs)
+
+    def preprocess(self, data: dict):
+        super().preprocess(data=data)
+
+        X = np.empty((data["X"].shape[0], data["X"].shape[1], data["X"].shape[2], data["X"].shape[3]), dtype=np.uint8)
+        Y = data["Y"]
+        for i in range(X.shape[0]):
+            for j in range(X.shape[3]):
+                X[i, :, :, j] = cv2.adaptiveThreshold(data["X"][i, :, :, j], 1, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
+                                                      cv2.THRESH_BINARY, 11, 2)
+
+        return dict(
+            X=X,
+            Y=Y,
+        )
diff --git a/tmu/composite/components/base.py b/tmu/composite/components/base.py
new file mode 100644
index 00000000..56a95c88
--- /dev/null
+++ b/tmu/composite/components/base.py
@@ -0,0 +1,67 @@
+import abc
+import numpy as np
+from pathlib import Path
+from typing import Union, Tuple
+import logging
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class TMComponent(abc.ABC):
+
+    def __init__(self, model_cls, model_config, epochs=1, **kwargs) -> None:
+        self.model_cls = model_cls
+        self.model_config = model_config
+        self.epochs = epochs
+
+        # Warn about unused kwargs
+        if kwargs:
+            print(f"Warning: unused keyword arguments: {kwargs}")
+
+        self.model_instance = self.model_cls(
+            **self.model_config.model_dump()
+        )
+
+    def model(self):
+        raise NotImplementedError
+
+    def preprocess(self, data: dict) -> dict:
+        assert "X" in data, "X must be in data"
+        assert "Y" in data, "Y must be in data"
+
+        # Check if this is called from a class that inherits from TMComponent
+        if not isinstance(self, TMComponent):
+            raise TypeError(f"{type(self).__name__} does not inherit from TMComponent")
+
+        return data
+
+    def fit(self, data: dict) -> None:
+        x_train, y_train = data["X"], data["Y"]
+
+        # Check if type is uint32
+        if x_train.dtype != np.uint32:
+            x_train = x_train.astype(np.uint32)
+
+        if y_train.dtype != np.uint32:
+            y_train = y_train.astype(np.uint32)
+
+        self.model_instance.fit(
+            x_train,
+            y_train,
+        )
+
+    def predict(self, data: dict) -> Tuple[np.array, np.array]:
+        X_test = data["X"]
+        return self.model_instance.predict(X_test, return_class_sums=True)
+
+    def save(self, path: Union[Path, str], format="pkl") -> None:
+        path = Path(path) if isinstance(path, str) else path
+
+        if format == "pkl":
+            import pickle
+            with open(path, "wb") as f:
+                pickle.dump(self, f)
+
+    def __str__(self):
+        params = '-'.join([f"{k}={v}" for k, v in self.model_config.model_dump().items()])
+        return f"{type(self).__name__}-{self.model_cls.__name__}-{params})"
diff --git a/tmu/composite/components/color_thermometer_scoring.py b/tmu/composite/components/color_thermometer_scoring.py
new file mode 100644
index 00000000..a3f50df6
--- /dev/null
+++ b/tmu/composite/components/color_thermometer_scoring.py
@@ -0,0 +1,26 @@
+import numpy as np
+from tmu.composite.components.base import TMComponent
+
+
+class ColorThermometerComponent(TMComponent):
+
+    def __init__(self, model_cls, model_config, resolution=8, **kwargs) -> None:
+        super().__init__(model_cls=model_cls, model_config=model_config, **kwargs)
+        self.resolution = resolution
+
+    def preprocess(self, data: dict):
+        super().preprocess(data=data)
+
+        X_org = data["X"]
+        Y = data["Y"]
+
+        X = np.empty((X_org.shape[0], X_org.shape[1], X_org.shape[2], X_org.shape[3], self.resolution), dtype=np.uint8)
+        for z in range(self.resolution):
+            X[:, :, :, :, z] = X_org[:, :, :, :] >= (z + 1) * 255 / (self.resolution + 1)
+
+        X = X.reshape((X_org.shape[0], X_org.shape[1], X_org.shape[2], 3 * self.resolution))
+
+        return dict(
+            X=X,
+            Y=Y,
+        )
diff --git a/tmu/composite/components/histogram_of_gradients.py b/tmu/composite/components/histogram_of_gradients.py
new file mode 100644
index 00000000..4ab94809
--- /dev/null
+++ b/tmu/composite/components/histogram_of_gradients.py
@@ -0,0 +1,72 @@
+import numpy as np
+import cv2
+
+from tmu.composite.components.base import TMComponent
+
+
+class HistogramOfGradientsComponent(TMComponent):
+
+    def __init__(
+            self,
+            model_cls,
+            model_config,
+            resolution=8,
+            block_size=12,
+            block_stride=4,
+            cell_size=4,
+            nbins=18,
+            deriv_aperture=1,
+            win_sigma=-1.,
+            histogram_norm_nype=0,
+            L2_hys_threshold=0.2,
+            gamma_correction=True,
+            n_levels=64,
+            signed_gradient=True,
+            **kwargs
+    ) -> None:
+        super().__init__(model_cls=model_cls, model_config=model_config, **kwargs)
+        self.resolution = resolution
+        self.block_size = block_size
+        self.block_stride = block_stride
+        self.cell_size = cell_size
+        self.nbins = nbins
+        self.deriv_aperture = deriv_aperture
+        self.win_sigma = win_sigma
+        self.histogram_norm_nype = histogram_norm_nype
+        self.L2_hys_threshold = L2_hys_threshold
+        self.gamma_correction = gamma_correction
+        self.n_levels = n_levels
+        self.signed_gradient = signed_gradient
+
+    def preprocess(self, data: dict):
+        super().preprocess(data=data)
+
+        X_org = data["X"]
+        Y = data["Y"]
+
+        win_size = X_org.shape[1]
+        hog = cv2.HOGDescriptor(
+            (win_size, win_size),
+            (self.block_size, self.block_size),
+            (self.block_stride, self.block_stride),
+            (self.cell_size, self.cell_size),
+            self.nbins,
+            self.deriv_aperture,
+            self.win_sigma,
+            self.histogram_norm_nype,
+            self.L2_hys_threshold,
+            self.gamma_correction,
+            self.n_levels,
+            self.signed_gradient
+        )
+
+        fd = hog.compute(X_org[0])
+        X = np.empty((X_org.shape[0], fd.shape[0]), dtype=np.uint32)
+        for i in range(X_org.shape[0]):
+            fd = hog.compute(X_org[i])
+            X[i] = fd >= 0.1
+
+        return dict(
+            X=X.astype(np.uint32),
+            Y=Y.astype(np.uint32)
+        )
diff --git a/tmu/composite/components/image/__init__.py b/tmu/composite/components/image/__init__.py
new file mode 100644
index 00000000..8a0136c9
--- /dev/null
+++ b/tmu/composite/components/image/__init__.py
@@ -0,0 +1,1 @@
+from tmu.composite.components.image.experimental import *
diff --git a/tmu/composite/components/image/experimental.py b/tmu/composite/components/image/experimental.py
new file mode 100644
index 00000000..8067edce
--- /dev/null
+++ b/tmu/composite/components/image/experimental.py
@@ -0,0 +1,520 @@
+import random
+
+from tmu.composite.components.base import TMComponent
+import cv2
+import numpy as np
+
+
+class RandomPatchExtractionComponent(TMComponent):
+    def preprocess(self, data: dict):
+        samples, rows, cols, ch = data["X"].shape
+        processed_X = np.empty((samples, 24, 24, ch), dtype=np.uint8)
+        patch_size = 24
+
+        for s in range(samples):
+            start_x = np.random.randint(0, rows - patch_size)
+            start_y = np.random.randint(0, cols - patch_size)
+            processed_X[s] = data["X"][s, start_x:start_x + patch_size, start_y:start_y + patch_size, :]
+
+        return {'X': processed_X, 'Y': data["Y"]}
+
+
+class SharpeningComponent(TMComponent):
+    def preprocess(self, data: dict):
+        kernel = np.array([[-1, -1, -1],
+                           [-1, 9, -1],
+                           [-1, -1, -1]])
+        samples, rows, cols, ch = data["X"].shape
+        processed_X = np.empty((samples, rows, cols, ch), dtype=np.uint8)
+
+        for s in range(samples):
+            processed_X[s] = cv2.filter2D(data["X"][s], -1, kernel)
+
+        return {'X': processed_X, 'Y': data["Y"]}
+
+
+class ChannelNormalizationComponent(TMComponent):
+    def preprocess(self, data: dict):
+        samples, rows, cols, ch = data["X"].shape
+        processed_X = np.empty((samples, rows, cols, ch), dtype=np.uint8)
+
+        for s in range(samples):
+            for i in range(3):  # assuming RGB channels
+                processed_X[s, ..., i] = (data["X"][s, ..., i] - np.mean(data["X"][s, ..., i])) / (
+                        np.std(data["X"][s, ..., i]) + 1e-7)
+
+        return {'X': processed_X, 'Y': data["Y"]}
+
+
+class SaltAndPepperNoiseComponent(TMComponent):
+    def preprocess(self, data: dict):
+        noise_amt = 0.02
+        samples, rows, cols, ch = data["X"].shape
+        processed_X = np.copy(data["X"])
+
+        for s in range(samples):
+            total_pixels = rows * cols
+
+            num_salt = np.ceil(noise_amt * total_pixels)
+            num_pepper = np.ceil(noise_amt * total_pixels)
+
+            coords_salt = [np.random.randint(0, i - 1, int(num_salt)) for i in [rows, cols]]
+            processed_X[s, coords_salt[0], coords_salt[1], :] = 255
+
+            coords_pepper = [np.random.randint(0, i - 1, int(num_pepper)) for i in [rows, cols]]
+            processed_X[s, coords_pepper[0], coords_pepper[1], :] = 0
+
+        return {'X': processed_X, 'Y': data["Y"]}
+
+
+class MorphologicalOperationsComponent(TMComponent):
+    def preprocess(self, data: dict):
+        samples, rows, cols, ch = data["X"].shape
+        processed_X = np.empty((samples, rows, cols, ch), dtype=np.uint8)
+        kernel = np.ones((5, 5), np.uint8)
+
+        for s in range(samples):
+            processed_X[s] = cv2.morphologyEx(data["X"][s], cv2.MORPH_CLOSE, kernel)
+
+        return {'X': processed_X, 'Y': data["Y"]}
+
+
+class InvertColorsComponent(TMComponent):
+    def preprocess(self, data: dict):
+        samples, rows, cols, ch = data["X"].shape
+        processed_X = np.empty((samples, rows, cols, ch), dtype=np.uint8)
+
+        for s in range(samples):
+            processed_X[s] = 255 - data["X"][s]
+
+        return {'X': processed_X, 'Y': data["Y"]}
+
+
+class GammaCorrectionComponent(TMComponent):
+    def preprocess(self, data: dict):
+        gamma = 2.0
+        invGamma = 1.0 / gamma
+        table = np.array([((i / 255.0) ** invGamma) * 255 for i in np.arange(0, 256)]).astype("uint8")
+
+        samples, rows, cols, ch = data["X"].shape
+        processed_X = np.empty((samples, rows, cols, ch), dtype=np.uint8)
+
+        for s in range(samples):
+            processed_X[s] = cv2.LUT(data["X"][s], table)
+
+        return {'X': processed_X, 'Y': data["Y"]}
+
+
+class OpticalFlowComponent(TMComponent):
+    def preprocess(self, data: dict):
+        samples, rows, cols, ch = data["X"].shape
+        processed_X = np.empty((samples, rows, cols, ch), dtype=np.uint8)
+
+        # Iterate through the batch
+        for s in range(samples):
+            # For the first image, use it as both current and previous to get zero flow.
+            prev_img = data["X"][s - 1] if s > 0 else data["X"][s]
+
+            flow = cv2.calcOpticalFlowFarneback(cv2.cvtColor(prev_img, cv2.COLOR_BGR2GRAY),
+                                                cv2.cvtColor(data["X"][s], cv2.COLOR_BGR2GRAY), None, 0.5, 3, 15, 3, 5,
+                                                1.2, 0)
+            hsv = np.zeros((rows, cols, 3), dtype=np.uint8)
+            hsv[..., 1] = 255
+            mag, ang = cv2.cartToPolar(flow[..., 0], flow[..., 1])
+            hsv[..., 0] = ang * 180 / np.pi / 2
+            hsv[..., 2] = cv2.normalize(mag, None, 0, 255, cv2.NORM_MINMAX)
+            processed_X[s] = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
+
+        return {'X': processed_X, 'Y': data["Y"]}
+
+
+class RandomRescaleComponent(TMComponent):
+    def preprocess(self, data: dict):
+        samples, rows, cols, ch = data["X"].shape
+        processed_X = np.empty((samples, rows, cols, ch), dtype=np.uint8)
+
+        for s in range(samples):
+            factor = random.uniform(0.8, 1.2)
+            rescaled = cv2.resize(data["X"][s], None, fx=factor, fy=factor, interpolation=cv2.INTER_LINEAR)
+
+            if factor > 1:  # Image is larger and needs to be cropped
+                start_x = random.randint(0, rescaled.shape[1] - cols)
+                start_y = random.randint(0, rescaled.shape[0] - rows)
+                rescaled = rescaled[start_y:start_y + rows, start_x:start_x + cols]
+            elif factor < 1:  # Image is smaller and needs to be padded
+                pad_x = (cols - rescaled.shape[1]) // 2
+                pad_y = (rows - rescaled.shape[0]) // 2
+                rescaled = cv2.copyMakeBorder(rescaled, pad_y, rows - rescaled.shape[0] - pad_y,
+                                              pad_x, cols - rescaled.shape[1] - pad_x,
+                                              cv2.BORDER_CONSTANT, value=[0, 0, 0])  # Using black for padding
+
+            processed_X[s] = rescaled
+
+        return {'X': processed_X, 'Y': data["Y"]}
+
+
+
+class CLAHEComponent:
+    def preprocess(self, data: dict):
+        # Extracting shape details
+        samples, rows, cols, ch = data["X"].shape
+
+        # Pre-allocation of the output array for efficiency
+        processed_X = np.empty((samples, rows, cols, ch), dtype=data["X"].dtype)
+
+        # Initializing CLAHE
+        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
+
+        # Applying CLAHE to each sample and each channel
+        for s in range(samples):
+            for i in range(ch):
+                processed_X[s, ..., i] = clahe.apply(data["X"][s, ..., i])
+
+        return {'X': processed_X, 'Y': data["Y"]}
+
+
+class RandomTranslationComponent(TMComponent):
+    def preprocess(self, data: dict):
+        samples, rows, cols, ch = data["X"].shape
+        processed_X = np.empty((samples, rows, cols, ch), dtype=np.uint8)
+
+        for s in range(samples):
+            tx, ty = random.randint(-10, 10), random.randint(-10, 10)
+            M = np.float32([[1, 0, tx], [0, 1, ty]])
+            processed_X[s] = cv2.warpAffine(data["X"][s], M, (cols, rows))
+
+        return {'X': processed_X, 'Y': data["Y"]}
+
+
+class RandomRotation90Component(TMComponent):
+    def preprocess(self, data: dict):
+        samples, rows, cols, ch = data["X"].shape
+        processed_X = np.empty((samples, rows, cols, ch), dtype=np.uint8)
+
+        for s in range(samples):
+            k = random.randint(0, 3)
+            processed_X[s] = np.rot90(data["X"][s], k)
+
+        return {'X': processed_X, 'Y': data["Y"]}
+
+
+class EdgeDetectionComponent(TMComponent):
+    def preprocess(self, data: dict):
+        samples, rows, cols, ch = data["X"].shape
+        processed_X = np.empty((samples, rows, cols, ch), dtype=np.uint8)
+
+        for s in range(samples):
+            edges = cv2.Canny(data["X"][s], 100, 200)
+            processed_X[s] = cv2.cvtColor(edges, cv2.COLOR_GRAY2BGR)
+
+        return {'X': processed_X, 'Y': data["Y"]}
+
+
+class GaborFilterComponent(TMComponent):
+    def preprocess(self, data: dict):
+        samples, rows, cols, ch = data["X"].shape
+        processed_X = np.empty((samples, rows, cols, ch), dtype=np.uint8)
+
+        kernel = cv2.getGaborKernel((21, 21), 5.0, np.pi / 4, 10.0, 0.5, 0, ktype=cv2.CV_32F)
+
+        for s in range(samples):
+            processed_X[s] = cv2.filter2D(data["X"][s], cv2.CV_8UC3, kernel)
+
+        return {'X': processed_X, 'Y': data["Y"]}
+
+
+class BilateralFilterComponent(TMComponent):
+    def preprocess(self, data: dict):
+        samples, rows, cols, ch = data["X"].shape
+        processed_X = np.empty((samples, rows, cols, ch), dtype=np.uint8)
+
+        for s in range(samples):
+            processed_X[s] = cv2.bilateralFilter(data["X"][s], 9, 75, 75)
+
+        return {'X': processed_X, 'Y': data["Y"]}
+
+
+class DCTComponent:
+    def preprocess(self, data: dict):
+        # Extracting shape details
+        samples, rows, cols, ch = data["X"].shape
+
+        # Pre-allocation of the output array for efficiency
+        processed_X = np.empty((samples, rows, cols, ch), dtype=data["X"].dtype)
+
+        # Applying DCT to each channel of each sample
+        for s in range(samples):
+            for c in range(ch):
+                processed_X[s, ..., c] = cv2.dct(data["X"][s, ..., c])
+
+        return {'X': processed_X, 'Y': data["Y"]}
+
+
+class HueSaturationValueAdjustmentComponent(TMComponent):
+    def preprocess(self, data: dict):
+        samples, rows, cols, ch = data["X"].shape
+        processed_X = np.empty((samples, rows, cols, ch), dtype=np.uint8)
+
+        for s in range(samples):
+            img = data["X"][s]
+            hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
+            hsv[..., 0] = (hsv[..., 0] + random.randint(-10, 10)) % 180
+            hsv[..., 1] = np.clip(hsv[..., 1] + random.randint(-40, 40), 0, 255)
+            hsv[..., 2] = np.clip(hsv[..., 2] + random.randint(-40, 40), 0, 255)
+            processed_X[s] = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
+
+        return {'X': processed_X, 'Y': data["Y"]}
+
+
+class ColorBalanceComponent(TMComponent):
+    def preprocess(self, data: dict):
+        samples, rows, cols, ch = data["X"].shape
+        processed_X = np.empty((samples, rows, cols, ch), dtype=np.uint8)
+
+        for s in range(samples):
+            img = data["X"][s]
+            for i in range(3):
+                img[..., i] = np.clip(img[..., i] + random.randint(-20, 20), 0, 255)
+            processed_X[s] = img
+
+        return {'X': processed_X, 'Y': data["Y"]}
+
+
+class GradientMagnitudeComponent(TMComponent):
+    def preprocess(self, data: dict):
+        samples, rows, cols, ch = data["X"].shape
+        processed_X = np.empty((samples, rows, cols, ch), dtype=np.uint8)
+
+        for s in range(samples):
+            img = data["X"][s]
+            sobelx = cv2.Sobel(img, cv2.CV_64F, 1, 0, ksize=3)
+            sobely = cv2.Sobel(img, cv2.CV_64F, 0, 1, ksize=3)
+            magnitude = np.sqrt(sobelx ** 2 + sobely ** 2)
+            processed_X[s] = cv2.normalize(magnitude, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
+
+        return {'X': processed_X, 'Y': data["Y"]}
+
+
+class FFTComponent(TMComponent):
+    def preprocess(self, data: dict):
+        samples, rows, cols, ch = data["X"].shape
+        processed_X = np.empty((samples, rows, cols, ch), dtype=np.uint8)
+
+        for s in range(samples):
+            img = data["X"][s]
+            f_transform = np.fft.fft2(img, axes=(0, 1))
+            f_transform_shifted = np.fft.fftshift(f_transform)
+            magnitude_spectrum = np.log(np.abs(f_transform_shifted) + 1)
+            processed_X[s] = cv2.normalize(magnitude_spectrum, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
+
+        return {'X': processed_X, 'Y': data["Y"]}
+
+
+class GaussianBlurComponent(TMComponent):
+    def preprocess(self, data: dict):
+        samples, rows, cols, ch = data["X"].shape
+        processed_X = np.empty((samples, rows, cols, ch), dtype=np.uint8)
+
+        for s in range(samples):
+            img = data["X"][s]
+            processed_X[s] = cv2.GaussianBlur(img, (5, 5), 0)
+
+        return {'X': processed_X, 'Y': data["Y"]}
+
+
+class ImageRotationComponent(TMComponent):
+    def preprocess(self, data: dict):
+        samples, rows, cols, ch = data["X"].shape
+        processed_X = np.empty((samples, rows, cols, ch), dtype=np.uint8)
+
+        for s in range(samples):
+            img = data["X"][s]
+            angle = random.uniform(0, 360)
+            M = cv2.getRotationMatrix2D((cols / 2, rows / 2), angle, 1)
+            processed_X[s] = cv2.warpAffine(img, M, (cols, rows))
+
+        return {'X': processed_X, 'Y': data["Y"]}
+
+
+
+
+class HorizontalFlipComponent(TMComponent):
+    def preprocess(self, data: dict):
+        samples, rows, cols, ch = data["X"].shape
+        processed_X = np.empty((samples, rows, cols, ch), dtype=np.uint8)
+
+        for s in range(samples):
+            img = data["X"][s]
+            processed_X[s] = cv2.flip(img, 1)
+
+        return {'X': processed_X, 'Y': data["Y"]}
+
+
+class VerticalFlipComponent(TMComponent):
+    def preprocess(self, data: dict):
+        samples, rows, cols, ch = data["X"].shape
+        processed_X = np.empty((samples, rows, cols, ch), dtype=np.uint8)
+
+        for s in range(samples):
+            img = data["X"][s]
+            processed_X[s] = cv2.flip(img, 0)
+
+        return {'X': processed_X, 'Y': data["Y"]}
+
+
+class BrightnessAdjustmentComponent(TMComponent):
+    def preprocess(self, data: dict):
+        samples, rows, cols, ch = data["X"].shape
+        processed_X = np.empty((samples, rows, cols, ch), dtype=np.uint8)
+
+        for s in range(samples):
+            img = data["X"][s]
+            factor = random.uniform(0.5, 1.5)
+            processed_X[s] = np.clip(img * factor, 0, 255).astype(np.uint8)
+
+        return {'X': processed_X, 'Y': data["Y"]}
+
+
+class ContrastAdjustmentComponent(TMComponent):
+    def preprocess(self, data: dict):
+        samples, rows, cols, ch = data["X"].shape
+        processed_X = np.empty((samples, rows, cols, ch), dtype=np.uint8)
+
+        for s in range(samples):
+            img = data["X"][s]
+            factor = random.uniform(0.5, 1.5)
+            processed_X[s] = np.clip(127.5 + factor * (img - 127.5), 0, 255).astype(np.uint8)
+
+        return {'X': processed_X, 'Y': data["Y"]}
+
+
+class ImageNormalizationComponent(TMComponent):
+    def preprocess(self, data: dict):
+        samples, rows, cols, ch = data["X"].shape
+        processed_X = np.empty((samples, rows, cols, ch), dtype=np.float32)
+
+        for s in range(samples):
+            img = data["X"][s]
+            processed_X[s] = cv2.normalize(img, None, 0, 1, cv2.NORM_MINMAX, dtype=cv2.CV_32F)
+
+        return {'X': processed_X, 'Y': data["Y"]}
+
+
+class ImageShiftComponent(TMComponent):
+    def preprocess(self, data: dict):
+        samples, rows, cols, ch = data["X"].shape
+        processed_X = np.empty((samples, rows, cols, ch), dtype=np.uint8)
+
+        for s in range(samples):
+            img = data["X"][s]
+            shift_x = random.randint(-10, 10)
+            shift_y = random.randint(-10, 10)
+            M = np.float32([[1, 0, shift_x], [0, 1, shift_y]])
+            processed_X[s] = cv2.warpAffine(img, M, (cols, rows))
+
+        return {'X': processed_X, 'Y': data["Y"]}
+
+
+class ChannelShuffleComponent(TMComponent):
+    def preprocess(self, data: dict):
+        samples, rows, cols, ch = data["X"].shape
+        processed_X = np.empty((samples, rows, cols, ch), dtype=np.uint8)
+
+        for s in range(samples):
+            img = data["X"][s]
+            channels = [0, 1, 2]
+            random.shuffle(channels)
+            processed_X[s] = img[..., channels]
+
+        return {'X': processed_X, 'Y': data["Y"]}
+
+
+class AffineTransformationComponent(TMComponent):
+
+    def preprocess(self, data: dict):
+        # Extracting shape details
+        samples, rows, cols, ch = data["X"].shape
+
+        # Pre-allocation of the output array for efficiency
+        processed_X = np.empty((samples, rows, cols, ch), dtype=data["X"].dtype)
+
+        # Defining transformation matrices
+        pts1 = np.float32([[50, 50], [200, 50], [50, 200]])
+        pts2 = np.float32([[10, 100], [200, 50], [100, 250]])
+        M = cv2.getAffineTransform(pts1, pts2)
+
+        # Applying the transformation to each sample
+        for i in range(samples):
+            processed_X[i] = cv2.warpAffine(data["X"][i], M, (cols, rows))
+
+        return {'X': processed_X, 'Y': data["Y"]}
+
+
+class PerspectiveTransformComponent(TMComponent):
+    def preprocess(self, data: dict):
+        samples, rows, cols, ch = data["X"].shape
+        processed_X = np.empty((samples, rows, cols, ch), dtype=np.uint8)
+
+        pts1 = np.float32([[56, 65], [368, 52], [28, 387], [389, 390]])
+        pts2 = np.float32([[0, 0], [300, 0], [0, 300], [300, 300]])
+        M = cv2.getPerspectiveTransform(pts1, pts2)
+
+        for s in range(samples):
+            img = data["X"][s]
+            processed_X[s] = cv2.warpPerspective(img, M, (cols, rows))
+
+        return {'X': processed_X, 'Y': data["Y"]}
+
+
+class NoiseInjectionComponent(TMComponent):
+    def preprocess(self, data: dict):
+        samples, rows, cols, ch = data["X"].shape
+        processed_X = np.empty((samples, rows, cols, ch), dtype=np.uint8)
+
+        for s in range(samples):
+            img = data["X"][s]
+            noise = np.random.randn(*img.shape) * 25  # Adjust noise level
+            processed_X[s] = np.clip(img + noise, 0, 255).astype(np.uint8)
+
+        return {'X': processed_X, 'Y': data["Y"]}
+
+
+class HistogramEqualizationComponent(TMComponent):
+    def preprocess(self, data: dict):
+        samples, rows, cols, ch = data["X"].shape
+        processed_X = np.empty((samples, rows, cols, ch), dtype=np.uint8)
+
+        for s in range(samples):
+            img = data["X"][s]
+            for i in range(3):
+                processed_X[s, ..., i] = cv2.equalizeHist(img[..., i])
+
+        return {'X': processed_X, 'Y': data["Y"]}
+
+
+class ColorJitterComponent:
+    def preprocess(self, data: dict):
+        # Extracting shape details
+        samples, rows, cols, ch = data["X"].shape
+
+        # Pre-allocation of the output array for efficiency
+        processed_X = np.empty((samples, rows, cols, ch), dtype=data["X"].dtype)
+
+        # Applying color jittering to each sample
+        for s in range(samples):
+            img = data["X"][s]
+
+            # Brightness, Contrast, Saturation adjustment
+            brightness = random.uniform(0.5, 1.5)
+            contrast = random.uniform(0.5, 1.5)
+            saturation = random.uniform(0.5, 1.5)
+
+            hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
+            hsv[..., 1] = np.clip(hsv[..., 1] * saturation, 0, 255)
+            img = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
+            img = np.clip(img * brightness, 0, 255).astype(np.uint8)
+            img = np.clip(127.5 + contrast * (img - 127.5), 0, 255).astype(np.uint8)
+
+            processed_X[s] = img
+
+        return {'X': processed_X, 'Y': data["Y"]}
diff --git a/tmu/composite/composite.py b/tmu/composite/composite.py
new file mode 100644
index 00000000..c942e1d1
--- /dev/null
+++ b/tmu/composite/composite.py
@@ -0,0 +1,249 @@
+import threading
+from collections import defaultdict
+from os import cpu_count
+from typing import Optional, Type, Union, List
+from pathlib import Path
+from multiprocessing import Pool, Manager
+import numpy as np
+from tqdm import tqdm
+
+from tmu.composite.callbacks.base import TMCompositeCallbackProxy, TMCompositeCallback
+from tmu.composite.components.base import TMComponent
+from tmu.composite.gating.base import BaseGate
+from tmu.composite.gating.linear_gate import LinearGate
+
+
+class TMCompositeBase:
+
+    def __init__(self, composite) -> None:
+        self.composite = composite
+
+    def _component_predict(self, component, data):
+        data_preprocessed = component.preprocess(data)
+        _, scores = component.predict(data_preprocessed)
+
+        votes = dict()
+        votes["composite"] = np.zeros_like(scores, dtype=np.float32)
+        votes[str(component)] = np.zeros_like(scores, dtype=np.float32)
+
+        for i in range(scores.shape[0]):
+            denominator = np.max(scores[i]) - np.min(scores[i])
+            score = 1.0 * scores[i] / denominator if denominator != 0 else 0
+            votes["composite"][i] += score
+            votes[str(component)][i] += score
+
+        return votes
+
+
+class TMCompositeMP(TMCompositeBase):
+
+    def __init__(self, composite) -> None:
+        super().__init__(composite=composite)
+
+    def _listener(self, queue, callbacks):
+        while True:
+            item = queue.get()
+            if item == 'DONE':
+                break
+            method, *args = item
+            for callback in callbacks:
+                getattr(callback, method)(*args)
+
+    @staticmethod
+    def _mp_fit(args: tuple) -> None:
+        idx, component, data_preprocessed, proxy_callback = args
+
+        if proxy_callback:
+            proxy_callback.on_train_composite_begin(composite=component)
+
+        epochs = component.epochs
+        pbar = tqdm(total=epochs, position=idx)
+        pbar.set_description(f"Component {idx}: {type(component).__name__}")
+        for epoch in range(epochs):
+            if proxy_callback:
+                proxy_callback.on_epoch_component_begin(component=component, epoch=epoch)
+            component.fit(data=data_preprocessed)
+            pbar.update(1)
+            if proxy_callback:
+                proxy_callback.on_epoch_component_end(component=component, epoch=epoch)
+
+        if proxy_callback:
+            proxy_callback.on_train_composite_end(composite=component)
+        return component
+
+    def fit(self, data: dict, callbacks: Optional[list[TMCompositeCallback]] = None) -> None:
+
+        with Manager() as manager:
+
+            if callbacks:
+                callback_queue = manager.Queue()  # Create a queue with the manager
+                callback_proxy = TMCompositeCallbackProxy(callback_queue)
+
+                # Start listener thread
+                listener_thread = threading.Thread(target=self._listener, args=(callback_queue, callbacks))
+                listener_thread.start()
+            else:
+                callback_proxy = None
+
+            with Pool() as pool:
+                data_preprocessed = [component.preprocess(data) for component in self.composite.components]
+                self.composite.components = pool.map(TMCompositeMP._mp_fit,
+                                                     ((idx, component, data_preprocessed[idx], callback_proxy) for
+                                                      idx, component in
+                                                      enumerate(self.composite.components)))
+
+            if callbacks:
+                callback_queue.put('DONE')  # Send done signal to listener
+                listener_thread.join()  # Wait for listener to process all logs
+
+    def predict(self, data: dict, votes: dict, gating_mask: np.ndarray) -> np.array:
+        # Determine number of processes based on available CPU cores
+        n_processes = min(cpu_count(), len(self.composite.components))
+        with Pool(n_processes) as pool:
+            results = pool.starmap(self._component_predict, [
+                (component, data) for i, component in enumerate(self.composite.components)
+            ])
+
+            # Aggregate results from each process
+            for i, result in enumerate(results):
+                for key, score in result.items():
+
+                    # Apply gating mask
+                    masked_score = score * gating_mask[:, i]
+
+                    if key not in votes:
+                        votes[key] = masked_score
+                    else:
+                        votes[key] += masked_score
+
+
+class TMCompositeSingleCPU(TMCompositeBase):
+
+    def __init__(self, composite) -> None:
+        super().__init__(composite=composite)
+
+    def fit(self, data: dict, callbacks: Optional[list[TMCompositeCallback]] = None) -> None:
+        data_preprocessed = [component.preprocess(data) for component in self.composite.components]
+        epochs_left = [component.epochs for component in self.composite.components]
+        pbars = [tqdm(total=component.epochs) for component in self.composite.components]
+        for idx, (pbar, component) in enumerate(zip(pbars, self.composite.components)):
+            pbar.set_description(f"Component {idx}: {type(component).__name__}")
+
+        [callback.on_train_composite_begin(composite=self) for callback in callbacks]
+        epoch = 0
+        while any(epochs_left):
+            for idx, component in enumerate(self.composite.components):
+                if epochs_left[idx] > 0:
+                    [callback.on_epoch_component_begin(component=component, epoch=epoch) for callback in callbacks]
+                    component.fit(data=data_preprocessed[idx])
+                    [callback.on_epoch_component_end(component=component, epoch=epoch) for callback in callbacks]
+                    pbars[idx].update(1)
+                    epochs_left[idx] -= 1
+
+            epoch += 1
+
+        [callback.on_train_composite_end(composite=self) for callback in callbacks]
+
+    def predict(self, data: dict, votes: dict, gating_mask: np.ndarray):
+        pbar = tqdm(total=len(self.composite.components))
+        for i, component in enumerate(self.composite.components):
+            pbar.set_description(f"Component {i}: {type(component).__name__}")
+            component_votes = self._component_predict(component, data)
+            for key, score in component_votes.items():
+
+                # Apply gating mask
+                masked_score = score * gating_mask[:, i]
+
+                if key not in votes:
+                    votes[key] = masked_score
+                else:
+                    votes[key] += masked_score
+            pbar.update(1)
+
+
+class TMComposite:
+
+    def __init__(
+            self,
+            components: Optional[list[TMComponent]] = None,
+            gate_function: Optional[Type[BaseGate]] = None,
+            gate_function_params: Optional[dict] = None,
+            use_multiprocessing: bool = False
+    ) -> None:
+        self.components: List[TMComponent] = components or []
+        self.use_multiprocessing = use_multiprocessing
+
+        if gate_function_params is None:
+            gate_function_params = dict()
+
+        self.gate_function_instance = gate_function(self, **gate_function_params) if gate_function else LinearGate(self,
+                                                                                                                   **gate_function_params)
+
+        self.logic = TMCompositeSingleCPU(composite=self) if not use_multiprocessing else TMCompositeMP(composite=self)
+
+    def fit(self, data: dict, callbacks: Optional[list[TMCompositeCallback]] = None) -> None:
+        self.logic.fit(data, callbacks)
+
+    def predict(self, data: dict) -> np.array:
+        votes = dict()
+
+        # Gating Mechanism
+        gating_mask: np.ndarray = self.gate_function_instance.predict(data)
+        assert gating_mask.shape[1] == len(self.components)
+        assert gating_mask.shape[0] == data["Y"].shape[0]
+
+        self.logic.predict(data, votes, gating_mask)
+
+        return {k: v.argmax(axis=1) for k, v in votes.items()}
+
+    def save_model(self, path: Union[Path, str], format="pkl") -> None:
+        path = Path(path) if isinstance(path, str) else path
+
+        if format == "pkl":
+            import pickle
+            with open(path, "wb") as f:
+                pickle.dump(self, f)
+        else:
+            raise NotImplementedError(f"Format {format} not supported")
+
+    def load_model(self, path: Union[Path, str], format="pkl") -> None:
+        path = Path(path) if isinstance(path, str) else path
+
+        if format == "pkl":
+            import pickle
+            with open(path, "rb") as f:
+                loaded_model = pickle.load(f)
+                self.__dict__.update(loaded_model.__dict__)
+        else:
+            raise NotImplementedError(f"Format {format} not supported")
+
+    def load_model_from_components(self, path: Union[Path, str], format="pkl") -> None:
+        path = Path(path) if isinstance(path, str) else path
+
+        if not path.is_dir():
+            raise ValueError(f"{path} is not a directory.")
+
+        # Get all files with the desired format
+        files = [f for f in path.iterdir() if f.is_file() and f.suffix == f".{format}"]
+
+        # Group files by component details
+        component_groups = defaultdict(list)
+        for file in files:
+            parts = file.stem.split('-')
+            epoch = int(parts[-1])
+            component_name = '-'.join(parts[:-1])
+            component_groups[component_name].append((epoch, file))
+
+        # Sort files in each group based on the epoch and get the highest epoch file
+        latest_files = [sorted(group, key=lambda x: x[0], reverse=True)[0][1] for group in component_groups.values()]
+
+        # Load all components and add to components list
+        self.components = []
+        if format == "pkl":
+            import pickle
+            for component_path in latest_files:
+                with open(component_path, "rb") as f:
+                    loaded_component = pickle.load(f)
+                    self.components.append(loaded_component)
+        else:
+            raise NotImplementedError(f"Format {format} not supported")
diff --git a/tmu/composite/config.py b/tmu/composite/config.py
new file mode 100644
index 00000000..ecb37780
--- /dev/null
+++ b/tmu/composite/config.py
@@ -0,0 +1,14 @@
+from pydantic import BaseModel, Field
+from typing import Union
+
+class TMBaseConfig(BaseModel):
+    pass
+
+class TMClassifierConfig(TMBaseConfig):
+    number_of_clauses: int = Field(10, description="Number of Clauses")
+    T: int = Field(500, description="T value")
+    s: float = Field(10.0, description="s value")
+    max_included_literals: int = Field(32, description="Max Included Literals")
+    platform: str = Field("GPU", description="Device to use")
+    weighted_clauses: bool = Field(True, description="Use Weighted Clauses or not")
+    patch_dim: Union[tuple, None] = Field(None, description="Size of Patch")
diff --git a/tmu/composite/gating/__init__.py b/tmu/composite/gating/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tmu/composite/gating/base.py b/tmu/composite/gating/base.py
new file mode 100644
index 00000000..9d448fec
--- /dev/null
+++ b/tmu/composite/gating/base.py
@@ -0,0 +1,14 @@
+
+class BaseGate:
+
+    def __init__(self, composite, **kwargs):
+        self.composite = composite
+
+    def preprocess(self, data: dict):
+        pass
+
+    def fit(self, data: dict) -> None:
+        pass
+
+    def predict(self, data: dict) -> dict:
+        pass
diff --git a/tmu/composite/gating/linear_gate.py b/tmu/composite/gating/linear_gate.py
new file mode 100644
index 00000000..45f066c0
--- /dev/null
+++ b/tmu/composite/gating/linear_gate.py
@@ -0,0 +1,16 @@
+import numpy as np
+from tmu.composite.gating.base import BaseGate
+
+
+class LinearGate(BaseGate):
+
+    def __init__(self, composite, **kwargs):
+        super().__init__(composite, **kwargs)
+
+    def predict(self, data: dict) -> np.ndarray:
+
+        # In the linear gate, we simply select all the components
+        n_components = len(self.composite.components)
+        n_data_points = data["Y"].shape[0]
+
+        return np.ones((n_data_points, n_components, 1), dtype=np.float32)
diff --git a/tmu/composite/gating/neural_gate.py b/tmu/composite/gating/neural_gate.py
new file mode 100644
index 00000000..b6b080e2
--- /dev/null
+++ b/tmu/composite/gating/neural_gate.py
@@ -0,0 +1,31 @@
+import numpy as np
+from tmu.composite.gating.base import BaseGate
+from keras.models import Sequential
+from keras.layers import Dense
+from keras.optimizers import Adam
+
+
+class NeuralGate(BaseGate):
+
+    def __init__(self, composite, input_dim, **kwargs):
+        super().__init__(composite, **kwargs)
+        self.model = self._build_model(input_dim)
+
+    def _build_model(self, input_dim):
+        """Construct a feed-forward neural network."""
+        model = Sequential()
+        model.add(Dense(128, input_dim=input_dim, activation='relu'))
+        model.add(Dense(64, activation='relu'))
+        model.add(Dense(len(self.composite.components), activation='softmax'))
+
+        model.compile(optimizer=Adam(), loss='categorical_crossentropy', metrics=['accuracy'])
+        return model
+
+    def train(self, X, y, epochs=10, batch_size=32):
+        """Train the neural network."""
+        self.model.fit(X, y, epochs=epochs, batch_size=batch_size)
+
+    def predict(self, data: dict) -> np.ndarray:
+        """Predict the gating mechanism."""
+        X_test = data["X"]
+        return self.model.predict(X_test)
diff --git a/tmu/composite/tuner.py b/tmu/composite/tuner.py
new file mode 100644
index 00000000..06a8739c
--- /dev/null
+++ b/tmu/composite/tuner.py
@@ -0,0 +1,144 @@
+import json
+import time
+import optuna
+from optuna.storages import JournalStorage, JournalFileStorage
+from joblib import Parallel, delayed
+import numpy as np
+
+from tmu.models.classification.vanilla_classifier import TMClassifier
+from tmu.composite.components.adaptive_thresholding import AdaptiveThresholdingComponent
+from tmu.composite.components.color_thermometer_scoring import ColorThermometerComponent
+from tmu.composite.components.histogram_of_gradients import HistogramOfGradientsComponent
+from tmu.composite.composite import TMComposite
+from tmu.composite.config import TMClassifierConfig
+
+
+class TMCompositeTuner:
+
+    def __init__(
+            self,
+            data_train,
+            data_test,
+            platform="CPU",
+            max_epochs=200,
+            n_jobs: int = 1,
+            callbacks=None,
+            use_multiprocessing=True,
+            study_name="TMComposite_study"
+    ):
+        self.data_train = data_train
+        self.data_test = data_test
+        self.last_accuracy = 0.0
+        self.n_components = 1
+        self.n_jobs = n_jobs
+        self.study_name = study_name
+        self.platform = platform
+        self.max_epochs = max_epochs
+        if callbacks is None:
+            callbacks = []
+
+        self.callbacks = callbacks
+        self.use_multiprocessing = use_multiprocessing
+
+    def objective(self, trial: optuna.trial.Trial) -> float:
+        components_list = []
+
+        for i in range(self.n_components):
+            component_type = trial.suggest_categorical(f'component_type_{i}',
+                                                       ['AdaptiveThresholdingComponent',
+                                                        'ColorThermometerComponent',
+                                                        'HistogramOfGradientsComponent'])
+
+            num_clauses = trial.suggest_int(f'num_clauses_{i}', 1000, 3000)
+            T = trial.suggest_int(f'T_{i}', 100, 1500)
+            s = trial.suggest_float(f's_{i}', 2.0, 15.0)
+            max_included_literals = trial.suggest_int(f'max_literals_{i}', 16, 64)
+            weighted_clauses = trial.suggest_categorical(f'weighted_clauses_{i}', [True, False])
+            epochs = trial.suggest_int(f'epochs_{i}', 1, self.max_epochs)
+
+            config = TMClassifierConfig(
+                num_clauses=num_clauses,
+                T=T,
+                s=s,
+                max_included_literals=max_included_literals,
+                platform=self.platform,
+                weighted_clauses=weighted_clauses
+            )
+
+            if component_type == 'AdaptiveThresholdingComponent':
+                patch_dim = (trial.suggest_int(f'patch_dim_1_{i}', 1, 10), trial.suggest_int(f'patch_dim_2_{i}', 1, 10))
+                config.patch_dim = patch_dim
+                components_list.append(AdaptiveThresholdingComponent(TMClassifier, config, epochs=epochs))
+
+            elif component_type == 'ColorThermometerComponent':
+                patch_dim = (trial.suggest_int(f'patch_dim_1_{i}', 1, 10), trial.suggest_int(f'patch_dim_2_{i}', 1, 10))
+                config.patch_dim = patch_dim
+                resolution = trial.suggest_int(f'resolution_{i}', 1, 10)
+                components_list.append(
+                    ColorThermometerComponent(TMClassifier, config, resolution=resolution, epochs=epochs))
+
+            elif component_type == 'HistogramOfGradientsComponent':
+                components_list.append(HistogramOfGradientsComponent(TMClassifier, config, epochs=epochs))
+
+        composite_model = TMComposite(components=components_list, use_multiprocessing=self.use_multiprocessing)
+
+        # Training and evaluation
+        composite_model.fit(
+            data=self.data_train,
+            callbacks=self.callbacks
+        )
+
+        preds = composite_model.predict(data=self.data_test)
+        accuracy = (preds['composite'] == self.data_test['Y'].flatten()).mean()
+
+        # Adjust number of components for next trial
+        if accuracy > self.last_accuracy:
+            self.n_components += 1
+        else:
+            self.n_components = max(1, self.n_components - 1)
+
+        self.last_accuracy = accuracy
+        return accuracy
+
+    def save_best_params(self, study, trial, filename="best_params.json"):
+        best_data = {
+            'params': study.best_params,
+            'value': trial.value
+        }
+        with open(filename, "w") as f:
+            json.dump(best_data, f)
+
+    def gradual_saving_callback(self, study, trial):
+        # Use np.isclose to handle potential floating-point precision issues
+        if np.isclose(trial.value, study.best_value, atol=1e-10):
+            self.save_best_params(study, trial, filename=f"best_params_trial_{trial.number}.json")
+
+    def retry_optimize(self, study, objective, n_trials, callbacks, max_retries=5, wait_time=2.0):
+        for _ in range(max_retries):
+            try:
+                study.optimize(objective, n_trials=n_trials, callbacks=callbacks)
+                return
+            except Exception as e:
+                if "database is locked" in str(e).lower():
+                    time.sleep(wait_time)
+                else:
+                    raise e
+        raise RuntimeError("Max retries reached for database access")
+
+    def tune(self, n_trials: int = 100):
+        storage = JournalStorage(JournalFileStorage("optuna-journal.log"))
+        with Parallel(n_jobs=self.n_jobs) as parallel:
+            if self.n_jobs == 1:
+                study = optuna.create_study(direction='maximize', pruner=optuna.pruners.MedianPruner(), storage=storage,
+                                            load_if_exists=True)
+                self.retry_optimize(study, self.objective, n_trials, [self.gradual_saving_callback])
+            else:
+                study = optuna.create_study(study_name=self.study_name, direction='maximize', storage=storage,
+                                            load_if_exists=True, pruner=optuna.pruners.MedianPruner())
+                parallel(
+                    delayed(self.retry_optimize)(study, self.objective, n_trials // self.n_jobs,
+                                                 [self.gradual_saving_callback])
+                    for i in range(self.n_jobs)
+                )
+
+        return study.best_params, study.best_value
diff --git a/tmu/data/cifar100.py b/tmu/data/cifar100.py
index 04490b13..53d9d8a9 100644
--- a/tmu/data/cifar100.py
+++ b/tmu/data/cifar100.py
@@ -2,7 +2,7 @@
 from typing import Dict
 from tmu.data import TMUDataset
 from tmu.data.utils.downloader import get_file
-
+from packaging.version import parse as parse_version
 
 class CIFAR100(TMUDataset):
     def _retrieve_dataset(self) -> Dict[str, np.ndarray]:
diff --git a/tmu/data/fashion_mnist.py b/tmu/data/fashion_mnist.py
index 29426b16..d1a7ccc3 100644
--- a/tmu/data/fashion_mnist.py
+++ b/tmu/data/fashion_mnist.py
@@ -5,6 +5,7 @@
 import sklearn
 from typing import Dict
 import numpy as np
+from packaging.version import parse as parse_version
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -17,9 +18,12 @@ def __init__(self):
 
     def _retrieve_dataset(self) -> Dict[str, np.ndarray]:
         kwargs = dict()
-        pyver = tuple([int(x) for x in sklearn.__version__.split(".")])
 
-        if pyver[0] >= 1 and pyver[1] >= 2:
+        # Parse the sklearn version string
+        sklearn_version = parse_version(sklearn.__version__)
+
+        # Check if the major version is >= 1 and the minor version is >= 2
+        if sklearn_version >= parse_version("1.2"):
             kwargs["parser"] = "pandas"
 
         X, y = fetch_openml(
@@ -51,9 +55,11 @@ def _transform(self, name, dataset):
 class KuzushijiMNIST(TMUDataset):
     def _retrieve_dataset(self) -> Dict[str, np.ndarray]:
         kwargs = dict()
-        pyver = tuple([int(x) for x in sklearn.__version__.split(".")])
+        # Parse the sklearn version string
+        sklearn_version = parse_version(sklearn.__version__)
 
-        if pyver[0] >= 1 and pyver[1] >= 2:
+        # Check if the major version is >= 1 and the minor version is >= 2
+        if sklearn_version >= parse_version("1.2"):
             kwargs["parser"] = "pandas"
 
         X, y = fetch_openml(
@@ -90,9 +96,11 @@ def __init__(self):
 
     def _retrieve_dataset(self) -> Dict[str, np.ndarray]:
         kwargs = dict()
-        pyver = tuple([int(x) for x in sklearn.__version__.split(".")])
+        # Parse the sklearn version string
+        sklearn_version = parse_version(sklearn.__version__)
 
-        if pyver[0] >= 1 and pyver[1] >= 2:
+        # Check if the major version is >= 1 and the minor version is >= 2
+        if sklearn_version >= parse_version("1.2"):
             kwargs["parser"] = "pandas"
 
         X, y = fetch_openml(
diff --git a/tmu/data/mnist.py b/tmu/data/mnist.py
index 645d7f26..887ffce4 100644
--- a/tmu/data/mnist.py
+++ b/tmu/data/mnist.py
@@ -31,6 +31,6 @@ def _retrieve_dataset(self) -> Dict[str, np.ndarray]:
 
     def _transform(self, name, dataset):
         if name.startswith("y"):
-            return dataset
+            return dataset.astype(np.uint32)
 
-        return np.where(dataset.reshape((dataset.shape[0], 28 * 28)) > 75, 1, 0)
+        return np.where(dataset.reshape((dataset.shape[0], 28 * 28)) > 75, 1, 0).astype(np.uint32)
diff --git a/tmu/data/tmu_datasource.py b/tmu/data/tmu_datasource.py
index f3ccfbf1..7bce675d 100644
--- a/tmu/data/tmu_datasource.py
+++ b/tmu/data/tmu_datasource.py
@@ -359,13 +359,18 @@ def get_dataset(
             Y_test = Y_test.values.flatten()
 
         if return_type == tuple:
-            return X_train, Y_train, X_test, Y_test
+            return (
+                X_train.astype(np.uint32),
+                Y_train.astype(np.uint32),
+                X_test.astype(np.uint32),
+                Y_test.astype(np.uint32)
+            )
         elif return_type == dict:
             return dict(
-                x_train=X_train,
-                y_train=Y_train,
-                x_test=X_test,
-                y_test=Y_test
+                x_train=X_train.astype(np.uint32),
+                y_train=Y_train.astype(np.uint32),
+                x_test=X_test.astype(np.uint32),
+                y_test=Y_test.astype(np.uint32)
             )
         else:
             raise RuntimeError("Invalid return_type. Should be set as 'dict' or 'tuple'")
diff --git a/tmu/experimental/__init__.py b/tmu/experimental/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tmu/experimental/models/__init__.py b/tmu/experimental/models/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tmu/models/attention.py b/tmu/experimental/models/attention.py
similarity index 100%
rename from tmu/models/attention.py
rename to tmu/experimental/models/attention.py
diff --git a/tmu/models/classification/multichannel_classifier.py b/tmu/experimental/models/multichannel_classifier.py
similarity index 100%
rename from tmu/models/classification/multichannel_classifier.py
rename to tmu/experimental/models/multichannel_classifier.py
diff --git a/tmu/models/classification/one_vs_one_classifier.py b/tmu/experimental/models/one_vs_one_classifier.py
similarity index 100%
rename from tmu/models/classification/one_vs_one_classifier.py
rename to tmu/experimental/models/one_vs_one_classifier.py
diff --git a/tmu/experimental/models/relational/__init__.py b/tmu/experimental/models/relational/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tmu/models/relational/vanilla_relational.py b/tmu/experimental/models/relational/vanilla_relational.py
similarity index 100%
rename from tmu/models/relational/vanilla_relational.py
rename to tmu/experimental/models/relational/vanilla_relational.py
diff --git a/tmu/lib/CMakeLists.txt b/tmu/lib/CMakeLists.txt
new file mode 100644
index 00000000..d7b1ffd8
--- /dev/null
+++ b/tmu/lib/CMakeLists.txt
@@ -0,0 +1,159 @@
+cmake_minimum_required(VERSION 3.15...3.27)
+project(tumlibpp)
+
+
+
+set(CMAKE_CXX_STANDARD 20)
+set(CMAKE_CXX_EXTENSIONS OFF) # Ensure -std=c++20 rather than -std=gnu++20
+
+option(BUILD_PYTHON "Build only the Python module" ON)
+option(BUILD_EXECUTABLE "Build the executable" ON)
+option(BUILD_STM32 "Build for STM32" OFF)
+
+
+
+include(ExternalProject)
+
+# Declare the span external project
+ExternalProject_Add(span
+        URL https://github.com/tcbrindle/span/archive/master.zip
+        #URL_HASH SHA256=<hash-value>
+        PREFIX ${CMAKE_BINARY_DIR}/_deps
+        CONFIGURE_COMMAND ""
+        BUILD_COMMAND ""
+        INSTALL_COMMAND ""
+        LOG_DOWNLOAD ON
+)
+
+# Declare the optional external project
+ExternalProject_Add(optional
+        URL https://github.com/TartanLlama/optional/archive/master.zip
+        #URL_HASH SHA256=<hash-value>
+        PREFIX ${CMAKE_BINARY_DIR}/_deps
+        CONFIGURE_COMMAND ""
+        BUILD_COMMAND ""
+        INSTALL_COMMAND ""
+        LOG_DOWNLOAD ON
+)
+set(span_SOURCE_DIR "${CMAKE_BINARY_DIR}/_deps/span-src")
+set(optional_SOURCE_DIR "${CMAKE_BINARY_DIR}/_deps/optional-src")
+
+
+
+
+
+# Optimize for Release build
+set(CMAKE_CONFIGURATION_TYPES "Release" CACHE STRING "" FORCE) # Focus on Release build
+
+message(STATUS "PYTHON: ${BUILD_PYTHON}")
+if(BUILD_PYTHON)
+
+    find_package(Python 3.8 COMPONENTS Interpreter Development.Module REQUIRED)
+
+    # Detect the installed nanobind package and import it into CMake
+    execute_process(
+            COMMAND "${Python_EXECUTABLE}" -m nanobind --cmake_dir
+            OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE NB_DIR)
+    list(APPEND CMAKE_PREFIX_PATH "${NB_DIR}")
+    find_package(nanobind CONFIG REQUIRED)
+endif()
+
+
+
+# High performance and optimization flags
+#set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -Ofast -ffast-math -march=native -DNDEBUG -flto -funroll-loops")
+
+add_library(
+        tmulib
+        STATIC
+        src/Attention.c
+        src/ClauseBank.c
+        src/ClauseBankSparse.c
+        src/WeightBank.c
+        src/Tools.c
+        src/random/pcg32_fast.c
+        src/random/xorshift128.c
+        cpp/include/utils/sparse_clause_container.h
+)
+target_include_directories(tmulib PUBLIC include)
+
+target_link_libraries(
+        tmulib
+        PUBLIC
+        $<$<CONFIG:Debug>: -fsanitize=address>
+)
+
+
+add_library(
+        tmulibpp
+        SHARED
+        cpp/src/memory.cpp
+)
+target_include_directories(
+        tmulibpp
+        PUBLIC
+        cpp/include
+        ${span_SOURCE_DIR}/include
+        ${optional_SOURCE_DIR}/include
+)
+target_link_libraries(
+        tmulibpp
+        PUBLIC
+        tmulib
+)
+
+
+
+
+IF(NOT BUILD_STM32)
+    target_compile_options(tmulibpp
+            PRIVATE
+            $<$<CONFIG:Release>:-Ofast -ffast-math -march=native -DNDEBUG -flto -funroll-loops >
+            $<$<CONFIG:Debug>:-O0 -g3 -DDEBUG -fsanitize=address>
+    )
+    target_compile_options(tmulib PRIVATE
+            $<$<CONFIG:Release>:-Ofast -ffast-math -march=native -DNDEBUG -flto -funroll-loops>
+            $<$<CONFIG:Debug>:-O0 -g3 -DDEBUG -fsanitize=address>
+    )
+ENDIF()
+
+
+
+
+IF(BUILD_EXECUTABLE)
+    add_executable(
+            tmulib_exe
+            cpp/main.cpp
+    )
+    target_link_libraries(tmulib_exe PRIVATE tmulibpp)
+
+    IF(NOT BUILD_STM32)
+        target_compile_options(tmulib_exe PRIVATE
+                $<$<CONFIG:Release>:-Ofast -ffast-math -march=native -DNDEBUG -flto>
+                $<$<CONFIG:Debug>:-O0 -g3 -DDEBUG -fsanitize=address>
+        )
+    ENDIF()
+ENDIF()
+
+
+
+IF(BUILD_PYTHON)
+    nanobind_add_module(
+            tmulibpy
+            NOMINSIZE
+            cpp/nb/tmulibcpp.cpp
+    )
+    install(TARGETS tmulibpy LIBRARY DESTINATION .)
+
+    target_link_libraries(tmulibpy PRIVATE tmulibpp)
+    target_compile_options(tmulibpy PRIVATE
+            $<$<CONFIG:Release>:-Os -ffast-math -march=native -DNDEBUG -flto>
+    )
+    set_target_properties(tmulibpy PROPERTIES
+            CXX_STANDARD 20
+            CXX_STANDARD_REQUIRED ON
+            CXX_EXTENSIONS OFF
+    )
+
+ENDIF()
+
diff --git a/tmu/lib/cpp/include/models/classifiers/tm_vanilla.h b/tmu/lib/cpp/include/models/classifiers/tm_vanilla.h
new file mode 100644
index 00000000..c5e6dbe1
--- /dev/null
+++ b/tmu/lib/cpp/include/models/classifiers/tm_vanilla.h
@@ -0,0 +1,689 @@
+//
+// Created by per on 3/2/24.
+//
+
+#ifndef TUMLIBPP_TM_VANILLA_H
+#define TUMLIBPP_TM_VANILLA_H
+#include <cmath>
+#include "utils/sparse_clause_container.h"
+#include "tm_weight_bank.h"
+#include "tm_clause_dense.h"
+#include "utils/tm_math.h"
+#include <tcb/span.hpp>
+#include <tl/optional.hpp>
+
+extern "C" {
+    #include "fast_rand.h"
+}
+
+template<class Type>
+class TMVanillaClassifier {
+
+public:
+
+
+    int T;
+    uint32_t number_of_clauses;
+    bool confidence_driven_updating;
+    tcb::span<int32_t> positive_clauses;
+    tcb::span<int32_t> negative_clauses;
+
+    // data views
+    std::vector<uint32_t> encoded_X_train_cached;
+    std::vector<uint32_t> encoded_X_test_vector;
+    std::vector<uint32_t> encoded_X_train_shape;
+    std::vector<uint32_t> encoded_X_test_shape;
+
+    bool weighted_clauses;
+    bool type_i_feedback;
+    bool type_ii_feedback;
+    bool type_iii_feedback;
+    bool type_i_p;
+    bool type_ii_p;
+    float type_i_ii_ratio;
+    int seed;
+    tl::optional<std::size_t> max_included_literals;
+    float s;
+    float d;
+    float clause_drop_p = 0.0; // TODO
+    float literal_drop_p = 0.0; // TODO
+    bool feature_negation = true; // TODO
+
+    bool boost_true_positive_feedback;
+    bool reuse_random_feedback;
+    tl::optional<std::vector<int>> patch_dim;
+    int32_t number_of_state_bits;
+    int32_t number_of_state_bits_ind;
+    int32_t batch_size;
+    bool incremental;
+
+
+    SparseClauseContainer<TMClauseBankDense<Type>> clause_banks;
+    SparseClauseContainer<TMWeightBank<Type>> weight_banks;
+    TMMemory<uint32_t> memory;
+
+
+    bool _is_initialized = false;
+
+    TMVanillaClassifier(
+            int _T,
+            float _s,
+            float _d,
+            uint32_t _number_of_clauses,
+            bool _confidence_driven_updating,
+            bool _weighted_clauses,
+            bool _type_i_feedback,
+            bool _type_ii_feedback,
+            bool _type_iii_feedback,
+            float _type_i_ii_ratio,
+            tl::optional<std::size_t> _max_included_literals,
+            bool _boost_true_positive_feedback,
+            bool _reuse_random_feedback,
+            tl::optional<std::vector<int>> _patch_dim,
+            int32_t _number_of_state_bits,
+            int32_t _number_of_state_bits_ind,
+            int32_t _batch_size,
+            bool _incremental,
+            int _seed
+    )
+    : T(_T)
+    , s(_s)
+    , d(_d)
+    , number_of_clauses(_number_of_clauses)
+    , confidence_driven_updating(_confidence_driven_updating)
+    , weighted_clauses(_weighted_clauses)
+    , type_i_feedback(_type_i_feedback)
+    , type_ii_feedback(_type_ii_feedback)
+    , type_iii_feedback(_type_iii_feedback)
+    , type_i_ii_ratio(_type_i_ii_ratio)
+    , max_included_literals(_max_included_literals)
+    , seed(_seed)
+    , clause_banks(_seed)
+    , weight_banks(_seed)
+    , boost_true_positive_feedback(_boost_true_positive_feedback)
+    , reuse_random_feedback(_reuse_random_feedback)
+    , patch_dim(_patch_dim)
+    , number_of_state_bits(_number_of_state_bits)
+    , number_of_state_bits_ind(_number_of_state_bits_ind)
+    , batch_size(_batch_size)
+    , incremental(_incremental)
+    , memory()
+
+
+    {
+        if(type_i_ii_ratio >= 1.0){
+            type_i_p = 1.0;
+            type_ii_p = 1.0 / type_i_ii_ratio;
+        }else{
+            type_i_p = type_i_ii_ratio;
+            type_ii_p = 1.0;
+        }
+
+    }
+
+    void initialize(std::vector<uint32_t> X_shape){
+        static_assert(sizeof(unsigned int) == sizeof(int32_t), "Sizes of unsigned int and int32_t must be the same.");
+
+        auto unsigned_positive_clauses = memory.getSegment(number_of_clauses);
+        auto unsigned_negative_clauses = memory.getSegment(number_of_clauses);
+        positive_clauses = tcb::span<int32_t>(reinterpret_cast<int32_t*>(unsigned_positive_clauses.data()), unsigned_positive_clauses.size());
+        negative_clauses = tcb::span<int32_t>(reinterpret_cast<int32_t*>(unsigned_negative_clauses.data()), unsigned_negative_clauses.size());
+    }
+
+
+    void init_after(){
+        // Fill the first half of positive_clauses with 1s
+        std::fill(positive_clauses.begin(), positive_clauses.begin() + number_of_clauses / 2, 1);
+
+        // Fill the second half of negative_clauses with 1s
+        std::fill(negative_clauses.begin() + number_of_clauses / 2, negative_clauses.end(), 1);
+
+    }
+
+    std::size_t get_required_memory_size(){
+        auto mem_size = 0;
+
+        // Negative clauses
+        mem_size += number_of_clauses;
+
+        // Positive clauses
+        mem_size += number_of_clauses;
+
+
+        return mem_size;
+    }
+
+    std::vector<uint32_t> mechanism_literal_active() {
+
+        auto item = *clause_banks.begin();
+        auto number_of_literals = item->number_of_literals;;
+
+        size_t number_of_ta_chunks = (number_of_literals + 31) / 32; // Calculate the number of 32-bit chunks
+        std::vector<uint32_t> literal_active(number_of_ta_chunks, 0); // Initialize with zeros
+
+        std::uniform_real_distribution<float> dist(0.0f, 1.0f); // Distribution for random floats between 0 and 1
+
+        // Iterate over each literal to determine its activation based on the drop probability
+        for (size_t k = 0; k < number_of_literals; ++k) {
+
+            auto rng = fast_rand() / static_cast<float>(UINT32_MAX);
+
+            if (rng >= literal_drop_p) { // Literal remains active
+                size_t ta_chunk = k / 32;
+                size_t chunk_pos = k % 32;
+                literal_active[ta_chunk] |= (1 << chunk_pos);
+            }
+        }
+
+        // If feature negation is not applied, clear the corresponding bits for the second half of literals
+        if (!feature_negation) {
+            for (size_t k = number_of_literals / 2; k < number_of_literals; ++k) {
+                size_t ta_chunk = k / 32;
+                size_t chunk_pos = k % 32;
+                literal_active[ta_chunk] &= ~(1 << chunk_pos);
+            }
+        }
+
+        return literal_active;
+    }
+
+
+
+    const std::vector<uint32_t> mechanism_clause_active() {
+        size_t weight_banks_size = weight_banks.size();
+        size_t total_elements = weight_banks_size * number_of_clauses;
+        std::vector<uint32_t> clause_active(total_elements, 0);
+
+        for (size_t idx = 0; idx < total_elements; ++idx) {
+            // Generate random float between 0 and 1
+            // generates a random number between 0 and 2^32 - 1
+            const auto random_value = fast_rand() / static_cast<float>(UINT32_MAX);
+
+            // Compare with clause_drop_p and set to 1 if the random value is greater or equal
+            clause_active[idx] = static_cast<uint32_t>(random_value >= clause_drop_p);
+        }
+
+        return clause_active;
+    }
+
+    float mechanism_compute_update_probabilities(bool is_target, int class_sum) {
+        // Confidence-driven updating method
+        if (confidence_driven_updating) {
+            return (T - std::abs(class_sum)) / T;
+        }
+
+        // Compute based on whether the class is the target or not
+        if (is_target) {
+            return (T - class_sum) / (2.0 * T);
+        }
+
+        return (T + class_sum) / (2.0 * T);
+    }
+
+
+
+    void mechanism_feedback(
+            bool is_target,
+            uint32_t target,
+            const tcb::span<Type>& clause_output,
+            float update_p,
+            const tcb::span<Type>& clause_active,
+            const tcb::span<Type>& literal_active,
+            const tcb::span<Type>& encoded_xi
+    ) {
+
+        const auto& clause_a = (is_target) ? positive_clauses : negative_clauses;
+        const auto& clause_b = (is_target) ? negative_clauses : positive_clauses;
+
+        const auto clause_active_target = clause_active.subspan(
+                target * number_of_clauses,
+                number_of_clauses
+        );
+
+        auto clause_active_elem_a = TMMath::elementwise_multiply<uint32_t, int32_t>(clause_active_target, clause_a);
+        auto clause_active_elem_b = TMMath::elementwise_multiply<uint32_t, int32_t>(clause_active_target, clause_b);
+
+        const auto clause_active_elem_a_span = tcb::span<uint32_t>(clause_active_elem_a.data(), clause_active_elem_a.size());
+        const auto clause_active_elem_b_span = tcb::span<uint32_t>(clause_active_elem_b.data(), clause_active_elem_b.size());
+
+
+        if(weighted_clauses){
+
+            if(is_target) {
+                weight_banks[target]->increment(
+                        clause_output, // clause_output
+                        update_p, // update_p
+                        clause_active, // clause_active
+                        false // positive_weights
+                );
+
+            }else{
+
+                weight_banks[target]->decrement(
+                        clause_output, // clause_output
+                        update_p, // update_p
+                        clause_active, // clause_active
+                        false // negative_weights
+                );
+
+            }
+
+
+        }
+
+
+        if(type_i_feedback){
+
+            clause_banks[target]->type_i_feedback(
+                    update_p * type_i_p, // update_p
+                    clause_active_elem_a_span, // clause_active
+                    literal_active, // literal_active
+                    encoded_xi // encoded_X_train
+            );
+        }
+
+        if(type_ii_feedback){
+            clause_banks[target]->type_ii_feedback(
+                    update_p * type_ii_p, // update_p
+                    clause_active_elem_b_span, // clause_active
+                    literal_active, // literal_active
+                    encoded_xi // encoded_X_train
+            );
+        }
+
+        if(type_iii_feedback){
+            clause_banks[target]->type_iii_feedback(
+                    update_p, // update_p
+                    clause_active_elem_a_span, // clause_active
+                    literal_active, // literal_active
+                    encoded_xi, // encoded_X_train
+                    true // target
+            );
+
+            clause_banks[target]->type_iii_feedback(
+                    update_p, // update_p
+                    clause_active_elem_b_span, // clause_active
+                    literal_active, // literal_active
+                    encoded_xi, // encoded_X_train
+                    false // target
+            );
+        }
+
+    }
+
+    int mechanism_clause_sum(
+            uint32_t target,
+            const tcb::span<Type>& clause_active,
+            const tcb::span<Type>& literal_active,
+            const tcb::span<Type>& encoded_xi
+    ){
+        auto clause_bank = clause_banks[target];
+
+        clause_bank->calculate_clause_outputs_update(
+                literal_active,
+                encoded_xi
+        );
+
+        const tcb::span<int32_t> clause_weights = weight_banks[target]->weights;
+
+        const tcb::span<uint32_t> clause_active_target = clause_active.subspan(
+                target * number_of_clauses,
+                number_of_clauses
+        );
+
+        const auto ca_mul_clause_weights = TMMath::elementwise_multiply<uint32_t , int32_t>(
+                clause_active_target,
+                clause_weights
+        );
+
+        int32_t class_sum = std::inner_product(
+                ca_mul_clause_weights.begin(), ca_mul_clause_weights.end(),
+                clause_bank->clause_output.begin(),
+                0 // Initial sum value
+        );
+
+        class_sum = TMMath::clamp(class_sum, -static_cast<int32_t>(T), static_cast<int32_t>(T));
+
+
+        return class_sum;
+    }
+
+    float _fit_sample_target(
+            int class_sum,
+            const tcb::span<Type>& clause_outputs,
+            bool is_target_class,
+            uint32_t target,
+            const tcb::span<Type>& clause_active,
+            const tcb::span<Type>& literal_active,
+            const tcb::span<Type>& encoded_xi
+    ){
+
+
+        auto update_p = mechanism_compute_update_probabilities(
+                is_target_class,
+                class_sum
+        );
+
+        mechanism_feedback(
+                is_target_class,
+                target,
+                clause_outputs,
+                update_p,
+                clause_active,
+                literal_active,
+                encoded_xi
+        );
+
+        return update_p;
+    }
+
+    void _fit_sample(
+        const tcb::span<Type>& clause_active,
+        const tcb::span<Type>& literal_active,
+        const tcb::span<Type>& encoded_xi,
+        uint32_t target,
+        tl::optional<uint32_t> not_target = tl::nullopt
+    ){
+        auto class_sum = mechanism_clause_sum(
+                target,
+                clause_active,
+                literal_active,
+                encoded_xi
+        );
+
+        const auto& clause_outputs = clause_banks[target]->clause_output;
+
+        float update_p_target = _fit_sample_target(
+                class_sum,
+                clause_outputs,
+                true,
+                target,
+                clause_active,
+                literal_active,
+                encoded_xi
+        );
+
+
+        if(!not_target.has_value()){
+            return; // TODO Dict
+        }
+
+        auto class_sum_not = mechanism_clause_sum(
+                not_target.value(),
+                clause_active,
+                literal_active,
+                encoded_xi
+        );
+        const auto clause_outputs_not = clause_banks[not_target.value()]->clause_output;
+
+        auto update_p_not_target = _fit_sample_target(
+                class_sum_not,
+                clause_outputs_not,
+                false,
+                not_target.value(),
+                clause_active,
+                literal_active,
+                encoded_xi
+        );
+
+        return; // TODO dict
+
+    }
+
+    void init(
+        const tcb::span<Type>& y,
+        const tcb::span<Type>& x,
+        const std::vector<int32_t>& X_shape
+    ){
+        if(_is_initialized){
+            return;
+        }
+        _is_initialized = true;
+
+
+        // Get unique classes from y (Set)
+        std::set<int> unique_classes(y.begin(), y.end());
+
+        // Convert to vector
+        std::vector<int> cls(unique_classes.begin(), unique_classes.end());
+
+
+        // Set up template functions
+        clause_banks.template_instance = std::make_shared<TMClauseBankDense<Type>>(
+                s,
+                d,
+                boost_true_positive_feedback,
+                reuse_random_feedback,
+                X_shape,
+                patch_dim,
+                max_included_literals,
+                number_of_clauses,
+                number_of_state_bits,
+                number_of_state_bits_ind,
+                batch_size,
+                incremental,
+                seed
+        );
+        clause_banks.populate(cls);
+
+        weight_banks.template_instance = std::make_shared<TMWeightBank<Type>>();
+        weight_banks.populate(cls);
+
+        /// Setup memory
+        auto num_classes = weight_banks.size();
+        auto mem_size = get_required_memory_size();
+        for(auto class_id : weight_banks.get_classes()){
+
+            auto weight_bank = weight_banks[class_id];
+            auto clause_bank = clause_banks[class_id];
+
+            mem_size += weight_bank->getRequiredMemorySize(number_of_clauses);
+            mem_size += clause_bank->getRequiredMemorySize();
+
+        }
+
+        mem_size += clause_banks.begin()->get()->getEncodedXiSize(X_shape);
+        mem_size += y.size();
+
+        memory.reserve(mem_size);
+
+        auto number_of_ta_chunks = clause_banks.begin()->get()->number_of_ta_chunks;
+        auto encoded_x_vector = clause_banks.begin()->get()->prepare_X(
+                x,
+                X_shape
+        );
+
+        // Retrieve memory segment for the dataset
+        //encoded_X_train_cached = memory.getSegment(encoded_x_vector.size()); // TODO not using cache
+        encoded_X_train_shape = std::vector<uint32_t>({
+            static_cast<uint32_t>(X_shape.at(0)),
+            static_cast<uint32_t>(number_of_ta_chunks)
+        });
+        encoded_X_train_cached = std::vector<uint32_t>(encoded_x_vector.size());
+
+        // Copy the encoded_x_vector to the memory
+        std::copy(encoded_x_vector.begin(), encoded_x_vector.end(), encoded_X_train_cached.begin());
+
+
+        for(auto class_id : weight_banks.get_classes()){
+            auto clause_bank = clause_banks[class_id];
+            auto weight_bank = weight_banks[class_id];
+            clause_bank->initialize(memory);
+            weight_bank->initialize(memory, number_of_clauses);
+        }
+
+        initialize(encoded_X_train_shape);
+        init_after();
+
+    }
+
+
+    void fit(
+            const tcb::span<Type>& y,
+            const tcb::span<Type>& x,
+            const std::vector<int32_t>& X_shape,
+            bool shuffle
+    ){
+
+        init(y, x,X_shape);
+
+
+        const auto encoded_X = tcb::span<Type>(encoded_X_train_cached.data(), encoded_X_train_cached.size());
+        const auto& encoded_X_shape = encoded_X_train_shape;
+        const auto num_features = encoded_X_shape.at(1);
+
+        auto clause_active_matrix = mechanism_clause_active();
+        const auto clause_active = tcb::span<uint32_t>(clause_active_matrix.data(), clause_active_matrix.size());
+
+        auto literal_active_vector = mechanism_literal_active();
+        const auto literal_active = tcb::span<uint32_t>(literal_active_vector.data(), literal_active_vector.size());
+
+        std::vector<int> sample_indices(encoded_X_shape.at(0));
+        TMMath::aRange(
+                encoded_X_shape.at(0),
+                shuffle,
+                sample_indices
+        );
+
+        for(auto i = 0; i < sample_indices.size(); i++){
+            const auto sample_idx = sample_indices[i];
+            const auto target = y[sample_idx];
+            const auto not_target = weight_banks.sample({target});
+
+
+            const auto encoded_xi = encoded_X.subspan(
+                    sample_idx * num_features,
+                    num_features
+            );
+
+            _fit_sample(
+                    clause_active,
+                    literal_active,
+                    encoded_xi,
+                    target,
+                    not_target
+            );
+        }
+
+    }
+
+    std::vector<int> predict_compute_class_sums(
+            const tcb::span<Type>& encoded_xi,
+            int sample_index,
+            int num_items,
+            bool clip_class_sum) {
+        std::vector<int> class_sums;
+        class_sums.reserve(weight_banks.size()); // Assuming weight_banks is accessible and defined.
+
+        for (const auto& class_id : weight_banks.get_classes()) { // Assuming get_classes() is a method returning class IDs.
+            const auto& weight_bank = weight_banks[class_id]; // Use const ref to avoid copies.
+            const auto& clause_bank = clause_banks[class_id]; // Use const ref to avoid copies.
+            const tcb::span<int32_t>& weights = weight_bank->weights; // Assuming weights is accessible and defined.
+
+            // Assuming calculate_clause_outputs_predict returns a tcb::span<uint32_t> and is correctly defined.
+            const tcb::span<uint32_t> clause_outputs = clause_bank->calculate_clause_outputs_predict(
+                    encoded_xi,
+                    sample_index,
+                    num_items
+            );
+
+            int class_sum = std::inner_product(
+                    clause_outputs.begin(),
+                    clause_outputs.begin() + std::min(clause_outputs.size(), weights.size()),
+                    weights.begin(),
+                    0);
+
+            if (clip_class_sum) {
+                // Assuming T is defined and TMMath::clamp is correctly implemented.
+                class_sum = TMMath::clamp(class_sum, -T, T);
+            }
+
+            class_sums.push_back(class_sum);
+        }
+
+        return class_sums;
+    }
+
+
+    const tcb::span<Type> getEncodedTestData(
+            const tcb::span<Type>& x,
+            const std::vector<int32_t>& X_shape
+    ){
+
+        if(encoded_X_test_vector.size() != 0){
+           return tcb::span<Type>(encoded_X_test_vector.data(), encoded_X_test_vector.size());
+        }
+
+        auto& clause_bank = *clause_banks.begin();
+        auto number_of_ta_chunks = clause_bank->number_of_ta_chunks;
+
+        encoded_X_test_vector = clause_bank->prepare_X(
+                x,
+                X_shape
+        );
+
+        encoded_X_test_shape = std::vector<uint32_t>(
+        {
+            static_cast<uint32_t>(X_shape.at(0)),
+            static_cast<uint32_t>(number_of_ta_chunks)
+        });
+
+        return tcb::span<Type>(encoded_X_test_vector.data(), encoded_X_test_vector.size());
+    }
+
+
+    const std::pair<std::vector<int>, tl::optional<std::vector<std::vector<int>>>> predict(
+            const tcb::span<Type>& X_test,
+            const std::vector<int32_t >& X_shape,
+            bool clip_class_sum = false,
+            bool return_class_sum = false) {
+
+
+        const auto encoded_X_test = getEncodedTestData(X_test, X_shape);
+        const auto num_items = encoded_X_test_shape.at(0);
+        const auto num_features = encoded_X_test_shape.at(1);
+
+        std::vector<int> argmax_indices(num_items, 0);
+        tl::optional<std::vector<std::vector<int>>> optional_class_sums;
+
+        if (return_class_sum) {
+            // Initialize class_sums_matrix only if needed
+            optional_class_sums.emplace(num_items, std::vector<int>(weight_banks.size(), 0));
+        }
+
+        for (int sample_index = 0; sample_index < num_items; ++sample_index) {
+            const auto encoded_xi = encoded_X_test.subspan(sample_index * num_features, num_features);
+
+            const auto class_sums_vector = predict_compute_class_sums(
+                    encoded_xi,
+                    sample_index,
+                    num_items,
+                    clip_class_sum
+            );
+
+            // Compute argmax
+            int argmax_index = std::distance(
+                    class_sums_vector.begin(),
+                    std::max_element(
+                            class_sums_vector.begin(),
+                            class_sums_vector.end()
+                    )
+            );
+            argmax_indices[sample_index] = argmax_index;
+
+            // Store the class sums if requested
+            if (return_class_sum) {
+                (*optional_class_sums)[sample_index] = std::move(class_sums_vector);
+            }
+        }
+
+        // Return both argmax_indices and optionally class_sums_matrix
+        return {
+            std::move(argmax_indices),
+            std::move(optional_class_sums)
+        };
+    }
+
+};
+
+#endif //TUMLIBPP_TM_VANILLA_H
diff --git a/tmu/lib/cpp/include/tm_clause_dense.h b/tmu/lib/cpp/include/tm_clause_dense.h
new file mode 100644
index 00000000..618f3f28
--- /dev/null
+++ b/tmu/lib/cpp/include/tm_clause_dense.h
@@ -0,0 +1,547 @@
+//
+// Created by per on 3/1/24.
+//
+
+#ifndef TUMLIBPP_TM_CLAUSE_DENSE_H
+#define TUMLIBPP_TM_CLAUSE_DENSE_H
+
+#include <cstddef>
+#include <vector>
+#include <stdexcept>
+#include <tuple>
+#include <tl/optional.hpp>
+#include <memory>
+#include <iostream>
+#include "tm_memory.h"
+
+extern "C" {
+    #include "ClauseBank.h"
+    #include "Tools.h"
+}
+
+template<class T>
+class TMClauseBankDense: public std::enable_shared_from_this<TMClauseBankDense<T>>{
+
+
+public:
+    float d;
+    float s;
+    bool boost_true_positive_feedback;
+    bool reuse_random_feedback;
+    std::size_t number_of_clauses;
+    std::size_t number_of_state_bits;
+    std::size_t number_of_state_bits_ind;
+    std::size_t batch_size;
+    std::size_t max_included_literals;
+    std::size_t number_of_features;
+    std::size_t number_of_literals;
+
+    std::size_t number_of_ta_chunks;
+    std::size_t number_of_patches;
+    std::tuple<int, int, int> dim;
+    std::tuple<int, int> patch_dim;
+
+
+    bool incremental_clause_evaluation_initialized = false;
+    bool incremental;
+
+
+    // Memory Segments
+    tcb::span<T> clause_output;
+    tcb::span<T> clause_output_batch;
+    tcb::span<T> clause_and_target;
+    tcb::span<T> clause_output_patchwise;
+    tcb::span<T> feedback_to_ta;
+    tcb::span<T> output_one_patches;
+    tcb::span<T> literal_clause_count;
+    tcb::span<T> type_ia_feedback_counter;
+    tcb::span<T> literal_clause_map;
+    tcb::span<T> literal_clause_map_pos;
+    tcb::span<T> false_literals_per_clause;
+    tcb::span<T> previous_xi;
+    tcb::span<T> clause_bank;
+    tcb::span<T> actions;
+    tcb::span<T> clause_bank_ind;
+
+private:
+
+    int seed;
+
+
+public:
+
+    TMClauseBankDense(
+        float _s,
+        float _d,
+        bool _boost_true_positive_feedback,
+        bool _reuse_random_feedback,
+        std::vector<int32_t> X_shape,
+        tl::optional<std::vector<int>> _patch_dim,
+        tl::optional<std::size_t> _max_included_literals,
+        std::size_t _number_of_clauses,
+        std::size_t _number_of_state_bits,
+        std::size_t _number_of_state_bits_ind,
+        std::size_t _batch_size,
+        bool _incremental,
+        int _seed = 0
+    )
+    : s(_s)
+    , d(_d)
+    , boost_true_positive_feedback(_boost_true_positive_feedback)
+    , reuse_random_feedback(_reuse_random_feedback)
+    , number_of_clauses(_number_of_clauses)
+    , number_of_state_bits(_number_of_state_bits)
+    , number_of_state_bits_ind(_number_of_state_bits_ind)
+    , batch_size(_batch_size)
+    , incremental(_incremental)
+    , seed(_seed)
+    {
+
+
+        // Validate and set the dimensions based on X_shape.
+        dim = TMClauseBankDense::getDim(X_shape);
+        patch_dim = TMClauseBankDense::getPatchDim(_patch_dim, dim);
+
+
+
+        // Calculate the number of features.
+        number_of_features = std::get<0>(patch_dim) * std::get<1>(patch_dim) * std::get<2>(dim) +
+                             (std::get<0>(dim) - std::get<0>(patch_dim)) +
+                             (std::get<1>(dim) - std::get<1>(patch_dim));
+
+        // Calculate the number of literals.
+        number_of_literals = number_of_features * 2;
+
+        // Calculate the number of patches.
+        number_of_patches = (std::get<0>(dim) - std::get<0>(patch_dim) + 1) *
+                            (std::get<1>(dim) - std::get<1>(patch_dim) + 1);
+
+        // Set maximum included literals or use the total number of literals.
+        max_included_literals = _max_included_literals.value_or(number_of_literals);
+        if(max_included_literals == 0 ){
+            max_included_literals = number_of_literals;
+        }
+
+        // Calculate the number of ternary association chunks.
+        number_of_ta_chunks = (number_of_literals - 1) / (sizeof(T) * 8) + 1; // Assuming
+    }
+
+    static std::tuple<int, int> getPatchDim(const tl::optional<std::vector<int>> _patch_dim, const std::tuple<int, int, int> dim){
+        // Set patch dimensions.
+        if(!_patch_dim){
+            return std::make_tuple(std::get<0>(dim) * std::get<1>(dim) * std::get<2>(dim), 1);
+        }else{
+            return std::make_tuple(_patch_dim.value()[0], _patch_dim.value()[1]);
+        }
+    }
+
+    static std::tuple<int, int, int> getDim(const std::vector<int32_t> X_shape){
+        if (X_shape.size() == 2) {
+            return {X_shape[1], 1, 1};
+        } else if (X_shape.size() == 3) {
+            return {X_shape[1], X_shape[2], 1};
+        } else if (X_shape.size() == 4) {
+            return {X_shape[1], X_shape[2], X_shape[3]};
+        } else {
+            throw std::invalid_argument("X_shape must be a 2D, 3D, or 4D tensor");
+        }
+    }
+
+
+    std::size_t getRequiredMemorySize() const {
+        return calculateTotalMemorySize();
+    }
+
+    bool initialize(TMMemory<T>& memory) {
+        // Allocate memory segments directly for each data structure.
+        clause_output = memory.getSegment(calculateClauseOutputSize());
+        clause_output_batch = memory.getSegment(calculateClauseOutputBatchSize());
+        clause_and_target = memory.getSegment(calculateClauseAndTargetSize());
+        clause_output_patchwise = memory.getSegment(calculateClauseOutputPatchwiseSize());
+        feedback_to_ta = memory.getSegment(calculateFeedbackToTaSize());
+        output_one_patches = memory.getSegment(calculateOutputOnePatchesSize());
+        literal_clause_count = memory.getSegment(calculateLiteralClauseCountSize());
+        type_ia_feedback_counter = memory.getSegment(calculateTypeIaFeedbackCounterSize());
+        literal_clause_map = memory.getSegment(calculateLiteralClauseMapSize());
+        literal_clause_map_pos = memory.getSegment(calculateLiteralClauseMapPosSize());
+        false_literals_per_clause = memory.getSegment(calculateFalseLiteralsPerClauseSize());
+        previous_xi = memory.getSegment(calculatePreviousXiSize());
+        clause_bank = memory.getSegment(calculateClauseBankSize());
+        actions = memory.getSegment(calculateActionsSize());
+        clause_bank_ind = memory.getSegment(calculateClauseTaChunksStateBitsIndSize());
+
+
+        initializeClauses();
+
+        return true; // Assuming all allocations are successful
+    }
+
+    void initializeClauses(){
+        // Set all bits to 1 except the last bit in each "chunk" of the 1D array
+        for (int i = 0; i < number_of_clauses; ++i) {
+            for (int j = 0; j < number_of_ta_chunks; ++j) {
+                for (int k = 0; k < number_of_state_bits - 1; ++k) { // Up to the second last
+                    int index = i * (number_of_ta_chunks * number_of_state_bits) + j * number_of_state_bits + k;
+                    clause_bank[index] = ~uint32_t(0); // Set all bits to 1
+                }
+                // Set the last bit to 0
+                int last_bit_index = i * (number_of_ta_chunks * number_of_state_bits) + j * number_of_state_bits + (number_of_state_bits - 1);
+                clause_bank[last_bit_index] = 0;
+            }
+        }
+
+        // Sett all bits to 1 for independent clauses
+        std::fill(clause_bank_ind.begin(), clause_bank_ind.end(), ~0);
+    }
+
+    std::size_t getEncodedXiSize(const std::vector<int32_t>& X_shape) const {
+        return X_shape.at(0) * number_of_patches * number_of_ta_chunks;
+    }
+
+    const std::vector<uint32_t> prepare_X(
+            const tcb::span<uint32_t>& x,
+            const std::vector<int32_t >& X_shape
+
+    ){
+        std::vector<uint32_t> encoded_X(X_shape.at(0) * number_of_patches * number_of_ta_chunks) ;
+
+        tmu_encode(
+                x.data(),
+                encoded_X.data(),
+                X_shape.at(0),
+                std::get<0>(dim),
+                std::get<1>(dim),
+                std::get<2>(dim),
+                std::get<0>(patch_dim),
+                std::get<1>(patch_dim),
+                1, // TODO
+                0 // TODO
+        );
+
+        return std::move(encoded_X);
+    }
+
+
+    // Assuming we have a method to calculate the position in clause_bank
+    size_t calculatePosition(size_t clause, size_t ta_chunk) const {
+        return clause * number_of_ta_chunks * number_of_state_bits + ta_chunk * number_of_state_bits;
+    }
+
+    void setTAState(size_t clause, size_t ta, unsigned int state) {
+        size_t ta_chunk = ta / 32;
+        size_t chunk_pos = ta % 32;
+        size_t pos = calculatePosition(clause, ta_chunk);
+        for (size_t b = 0; b < number_of_state_bits; ++b) {
+            if (state & (1 << b)) {
+                clause_bank[pos + b] |= (1 << chunk_pos);
+            } else {
+                clause_bank[pos + b] &= ~(1 << chunk_pos);
+            }
+        }
+    }
+
+    unsigned int getTAState(size_t clause, size_t ta) {
+        size_t ta_chunk = ta / 32;
+        size_t chunk_pos = ta % 32;
+        size_t pos = calculatePosition(clause, ta_chunk);
+        unsigned int state = 0;
+        for (size_t b = 0; b < number_of_state_bits; ++b) {
+            if (clause_bank[pos + b] & (1 << chunk_pos)) {
+                state |= (1 << b);
+            }
+        }
+        return state;
+    }
+
+    bool get_ta_action(size_t clause, size_t ta) {
+        size_t ta_chunk = ta / 32;
+        size_t chunk_pos = ta % 32;
+        size_t pos = calculatePosition(clause, ta_chunk) + number_of_state_bits - 1;
+        return (clause_bank[pos] & (1 << chunk_pos)) > 0;
+    }
+
+    tcb::span<T> calculateClauseOutputsPredict() {
+        return clause_output;
+    }
+
+    std::size_t calculateClauseBankSize() const {
+        return number_of_clauses * number_of_ta_chunks * number_of_state_bits;
+    }
+
+    std::size_t calculateClauseOutputBatchSize() const {
+        return number_of_clauses * batch_size;
+    }
+
+
+
+
+    void type_i_feedback(
+            float update_p,
+            const tcb::span<T>& clause_active,
+            const tcb::span<T>& literal_active,
+            const tcb::span<T>& encoded_xi
+    ){
+
+        cb_type_i_feedback(
+                clause_bank.data(),
+                feedback_to_ta.data(),
+                output_one_patches.data(),
+                number_of_clauses,
+                number_of_literals,
+                number_of_state_bits,
+                number_of_patches,
+                update_p,
+                s,
+                boost_true_positive_feedback,
+                reuse_random_feedback,
+                max_included_literals,
+                clause_active.data(),
+                literal_active.data(),
+                encoded_xi.data()
+        );
+
+        incremental_clause_evaluation_initialized = false;
+
+    }
+
+    void type_ii_feedback(
+            float update_p,
+            const tcb::span<T>& clause_active,
+            const tcb::span<T>& literal_active,
+            const tcb::span<T>& encoded_xi
+    ){
+        cb_type_ii_feedback(
+            clause_bank.data(),
+            output_one_patches.data(),
+            number_of_clauses,
+            number_of_literals,
+            number_of_state_bits,
+            number_of_patches,
+            update_p,
+            clause_active.data(),
+            literal_active.data(),
+            encoded_xi.data()
+        );
+
+        incremental_clause_evaluation_initialized = false;
+    }
+
+    void type_iii_feedback(
+            float update_p,
+            const tcb::span<T>& clause_active,
+            const tcb::span<T>& literal_active,
+            const tcb::span<T>& encoded_X_train,
+            bool target
+    ){
+        cb_type_iii_feedback(
+                clause_bank.data(),
+                clause_bank_ind.data(),
+                clause_and_target.data(),
+                output_one_patches.data(),
+                number_of_clauses,
+                number_of_literals,
+                number_of_state_bits,
+                number_of_state_bits_ind,
+                number_of_patches,
+                update_p,
+                d,
+                clause_active.data(),
+                literal_active.data(),
+                encoded_X_train.data(), // TODO
+                target
+        );
+
+        incremental_clause_evaluation_initialized = false;
+    }
+
+    void calculate_clause_outputs_update(
+            const tcb::span<T>& literal_active,
+            const tcb::span<T>& encoded_xi
+    ){
+        cb_calculate_clause_outputs_update(
+                clause_bank.data(),
+                number_of_clauses,
+                number_of_literals,
+                number_of_state_bits,
+                number_of_patches,
+                clause_output.data(),
+                literal_active.data(),
+                encoded_xi.data()
+        );
+
+    }
+    
+    
+    const tcb::span<T> calculate_clause_outputs_predict(
+            const tcb::span<T>& encoded_xi,
+            std::size_t sample_index,
+            std::size_t n_items
+    ){
+
+        if(!incremental){
+            cb_calculate_clause_outputs_predict(
+                    clause_bank.data(),
+                    number_of_clauses,
+                    number_of_literals,
+                    number_of_state_bits,
+                    number_of_patches,
+                    clause_output.data(),
+                    encoded_xi.data()
+            );
+
+            return clause_output;
+        }
+
+        if(!incremental_clause_evaluation_initialized){
+
+            cb_initialize_incremental_clause_calculation(
+                    clause_bank.data(),
+                    literal_clause_map.data(),
+                    literal_clause_map_pos.data(),
+                    false_literals_per_clause.data(),
+                    number_of_clauses,
+                    number_of_literals,
+                    number_of_state_bits,
+                    previous_xi.data()
+            );
+
+            incremental_clause_evaluation_initialized = true;
+        }
+
+        if(sample_index % batch_size == 0){
+            cb_calculate_clause_outputs_incremental_batch(
+                    literal_clause_map.data(),
+                    literal_clause_map_pos.data(),
+                    false_literals_per_clause.data(),
+                    number_of_clauses,
+                    number_of_literals,
+                    number_of_patches,
+                    clause_output_batch.data(),
+                    previous_xi.data(),
+                    encoded_xi.data(),
+                    std::min(batch_size, n_items - sample_index)
+
+            );
+        }
+
+        size_t start_index = (sample_index % batch_size) * number_of_clauses;
+
+        // Create a span that represents the slice of the clause outputs for the entry 'e'
+        tcb::span<T> cl(
+                clause_output_batch.data() + start_index,
+                number_of_clauses
+        );
+
+        return cl;
+    }
+
+
+private:
+    std::size_t calculateClauseOutputSize() const {
+        return number_of_clauses;
+    }
+
+    std::size_t calculateClauseAndTargetSize() const {
+        return number_of_clauses * number_of_ta_chunks;
+    }
+
+    std::size_t calculateClauseOutputPatchwiseSize() const {
+        return number_of_clauses * number_of_patches;
+    }
+
+    std::size_t calculateFeedbackToTaSize() const {
+        return number_of_ta_chunks;
+    }
+
+    std::size_t calculateOutputOnePatchesSize() const {
+        return number_of_patches;
+    }
+
+    std::size_t calculateLiteralClauseCountSize() const {
+        return number_of_literals;
+    }
+
+    std::size_t calculateTypeIaFeedbackCounterSize() const {
+        return number_of_clauses;
+    }
+
+    std::size_t calculateLiteralClauseMapSize() const {
+        return number_of_literals * number_of_clauses;
+    }
+
+    std::size_t calculateLiteralClauseMapPosSize() const {
+        return number_of_literals;
+    }
+
+    std::size_t calculateFalseLiteralsPerClauseSize() const {
+        return number_of_clauses * number_of_patches;
+    }
+
+    std::size_t calculatePreviousXiSize() const {
+        return number_of_ta_chunks * number_of_patches;
+    }
+
+    std::size_t calculateActionsSize() const {
+        return number_of_ta_chunks;
+    }
+
+    std::size_t calculateClauseTaChunksStateBitsIndSize() const {
+        return number_of_clauses * number_of_ta_chunks * number_of_state_bits_ind;
+    }
+
+    std::size_t calculateTotalMemorySize() const {
+        return
+            calculateClauseOutputSize() +
+            calculateClauseOutputBatchSize() +
+            calculateClauseAndTargetSize() +
+            calculateClauseOutputPatchwiseSize() +
+            calculateFeedbackToTaSize() +
+            calculateOutputOnePatchesSize() +
+            calculateLiteralClauseCountSize() +
+            calculateTypeIaFeedbackCounterSize() +
+            calculateLiteralClauseMapSize() +
+            calculateLiteralClauseMapPosSize() +
+            calculateFalseLiteralsPerClauseSize() +
+            calculatePreviousXiSize() +
+            calculateClauseBankSize() +
+            calculateActionsSize() +
+            calculateClauseTaChunksStateBitsIndSize();
+    }
+
+    void calculate_literal_clause_frequency(){
+        // clause_active
+        // TODO
+    }
+
+    void calculate_clause_outputs_predict(){
+        // TODO
+    }
+
+    void included_literals(){
+        // TODO
+    }
+
+    void get_literals(){
+        // TODO
+    }
+
+    void number_of_include_actions(){
+        // TODO
+    }
+
+    void prepare_Xi(){
+        // TODO
+    }
+
+    void prepare_X_autoencoder(){
+        // TODO
+    }
+
+    void produce_autoencoder_example(){
+        // TODO
+    }
+
+};
+
+
+#endif //TUMLIBPP_TM_CLAUSE_DENSE_H
diff --git a/tmu/lib/cpp/include/tm_memory.h b/tmu/lib/cpp/include/tm_memory.h
new file mode 100644
index 00000000..cabb4606
--- /dev/null
+++ b/tmu/lib/cpp/include/tm_memory.h
@@ -0,0 +1,88 @@
+#ifndef TMU_MEMORY_H
+#define TMU_MEMORY_H
+
+#include <vector>
+#include <functional>
+#include <cstdint>
+#include <stdexcept>
+#include <tcb/span.hpp>
+
+template<typename T>
+class TMMemory {
+public:
+    using ReallocateCallback = std::function<void(std::size_t oldCapacity, std::size_t newCapacity)>;
+
+private:
+    alignas(64) std::vector<T> memory;
+    ReallocateCallback reallocateCallback;
+    std::size_t capacity;
+    std::size_t cursor = 0; // Cursor to track the current position
+    std::size_t _preallocate_size = 0; // used to determine how much memory to allocate before the first use
+
+public:
+    TMMemory() : capacity(0) {}
+
+    TMMemory(const std::vector<T>& data, std::size_t _cursor)
+    : memory(data)
+    , capacity(data.capacity())
+    , cursor(_cursor)
+    {}
+
+    void setReallocateCallback(ReallocateCallback callback) {
+        reallocateCallback = std::move(callback);
+    }
+
+    void addMemory(std::size_t size) {
+        _preallocate_size += size;
+    }
+
+    void reserve(std::size_t newCapacity) {
+        if (newCapacity > memory.capacity()) {
+            std::size_t oldCapacity = memory.capacity();
+            memory.reserve(_preallocate_size + newCapacity);
+
+            // Fill newly allocated memory with zeros
+            memory.resize(newCapacity);
+            std::fill(memory.begin() + oldCapacity, memory.end(), 0);
+
+            capacity = memory.capacity();
+            if (reallocateCallback) {
+                reallocateCallback(oldCapacity, capacity);
+            }
+        }
+    }
+
+    void push_back(T value) {
+        if (memory.capacity() == memory.size()) {
+            std::size_t oldCapacity = memory.capacity();
+            memory.push_back(value);
+            std::size_t newCapacity = memory.capacity();
+            if (oldCapacity != newCapacity && reallocateCallback) {
+                reallocateCallback(oldCapacity, newCapacity);
+            }
+        } else {
+            memory.push_back(value);
+        }
+    }
+
+    tcb::span<T> getSegment(std::size_t segmentSize){
+        if (cursor + segmentSize > memory.size()) {
+            throw std::out_of_range("Segment exceeds memory bounds. Requested '" + std::to_string(cursor + segmentSize) + " bytes' is more than '" + std::to_string(memory.size() - cursor) + "' bytes.");
+        }
+
+        tcb::span<T> segmentView(&memory[cursor], segmentSize);
+        cursor += segmentSize; // Advance the cursor
+        return segmentView;
+    }
+
+    // Resets the cursor to the beginning or to a specified position
+    void resetCursor(std::size_t newCursor = 0) {
+        if (newCursor > memory.size()) {
+            throw std::out_of_range("New cursor position exceeds memory bounds.");
+        }
+        cursor = newCursor;
+    }
+
+};
+
+#endif //TMU_MEMORY_H
diff --git a/tmu/lib/cpp/include/tm_weight_bank.h b/tmu/lib/cpp/include/tm_weight_bank.h
new file mode 100644
index 00000000..9d2ac96b
--- /dev/null
+++ b/tmu/lib/cpp/include/tm_weight_bank.h
@@ -0,0 +1,102 @@
+//
+// Created by per on 3/2/24.
+//
+
+#ifndef TUMLIBPP_TM_WEIGHT_BANK_H
+#define TUMLIBPP_TM_WEIGHT_BANK_H
+
+#include <span>
+#include "tm_memory.h"
+#include "tm_clause_dense.h"
+#include <assert.h>
+#include <iostream>
+extern "C" {
+    #include "WeightBank.h"
+}
+
+class TMWeightBankPresets {
+
+
+
+public:
+    static tcb::span<int32_t> positive_one_negative_minus_one(tcb::span<int32_t>& weights, std::size_t number_of_clauses) {
+        assert (weights.size() == number_of_clauses);
+
+        // Fill weights with 1 and -1
+        for (int i = 0; i < number_of_clauses / 2; i++) {
+            weights[i] = 1;
+            weights[i + (number_of_clauses / 2)] = -1;
+        }
+
+        return weights;
+    }
+
+};
+
+
+template<class T>
+class TMWeightBank: public std::enable_shared_from_this<TMWeightBank<T>> {
+
+
+
+public:
+    tcb::span<int32_t> weights;
+
+    TMWeightBank(){};
+
+    void initialize(TMMemory<T>& memory, std::size_t number_of_clauses) {
+
+        tcb::span<uint32_t> originalSpan = memory.getSegment(number_of_clauses);
+        tcb::span<int32_t> reinterpretedSpan(reinterpret_cast<int32_t*>(originalSpan.data()), originalSpan.size());
+        weights = reinterpretedSpan;
+
+        TMWeightBankPresets::positive_one_negative_minus_one(weights, number_of_clauses);
+    }
+
+    std::size_t getRequiredMemorySize(std::size_t number_of_clauses) {
+        return number_of_clauses;
+    }
+
+    void increment(
+            const tcb::span<uint32_t> clause_output,
+            float update_p,
+            const tcb::span<uint32_t>& clause_active,
+            bool positive_weights
+    ){
+
+        wb_increment(
+                weights.data(),
+                weights.size(),
+            clause_output.data(),
+            update_p,
+            clause_active.data(),
+            positive_weights
+        );
+    }
+
+    void decrement(
+            const tcb::span<uint32_t> clause_output,
+            float update_p,
+            const tcb::span<uint32_t>& clause_active,
+            bool negative_weights
+    ){
+
+        wb_decrement(
+            weights.data(),
+            weights.size(),
+            clause_output.data(),
+            update_p,
+            clause_active.data(),
+            negative_weights
+        );
+    }
+
+    const tcb::span<int32_t> getWeights() {
+        return weights;
+    }
+
+
+
+};
+
+#endif //TUMLIBPP_TM_WEIGHT_BANK_H
diff --git a/tmu/lib/cpp/include/utils/sparse_clause_container.h b/tmu/lib/cpp/include/utils/sparse_clause_container.h
new file mode 100644
index 00000000..1c3bdd32
--- /dev/null
+++ b/tmu/lib/cpp/include/utils/sparse_clause_container.h
@@ -0,0 +1,152 @@
+//
+// Created by per on 3/2/24.
+//
+
+#ifndef TUMLIBPP_SPARSE_CLAUSE_CONTAINER_H
+#define TUMLIBPP_SPARSE_CLAUSE_CONTAINER_H
+#include <utility>
+#include <vector>
+#include <unordered_map>
+#include <random>
+#include <set>
+#include <memory>
+#include <functional>
+#include <iostream>
+#include <tl/optional.hpp>
+
+
+#include "tm_clause_dense.h"
+
+extern "C" {
+    #include "fast_rand.h"
+}
+
+
+template<typename ClauseType>
+class ValueIterator : public std::iterator<std::forward_iterator_tag, std::shared_ptr<ClauseType>> {
+private:
+    using MapIterator = typename std::unordered_map<int, std::shared_ptr<ClauseType>>::iterator;
+    MapIterator it;
+
+public:
+    explicit ValueIterator(MapIterator it) : it(it) {}
+
+    ValueIterator& operator++() { ++it; return *this; } // Prefix increment
+    ValueIterator operator++(int) { ValueIterator tmp = *this; ++(*this); return tmp; } // Postfix increment
+
+    bool operator==(const ValueIterator& other) const { return it == other.it; }
+    bool operator!=(const ValueIterator& other) const { return it != other.it; }
+
+    std::shared_ptr<ClauseType>& operator*() const { return it->second; }
+    std::shared_ptr<ClauseType>* operator->() const { return &it->second; }
+};
+
+
+template<typename Iterator>
+class IteratorRange {
+private:
+    Iterator beginIterator, endIterator;
+
+public:
+    IteratorRange(Iterator begin, Iterator end)
+            : beginIterator(begin), endIterator(end) {}
+
+    Iterator begin() const { return beginIterator; }
+    Iterator end() const { return endIterator; }
+};
+
+template<class ClauseType>
+class SparseClauseContainer {
+    std::mt19937 rng;
+    std::vector<int> classes;
+    std::unordered_map<int, std::shared_ptr<ClauseType>> d;
+
+public:
+    std::shared_ptr<ClauseType> template_instance;
+
+
+    explicit SparseClauseContainer(unsigned int random_seed)
+    : rng(random_seed)
+    {
+    }
+
+    ~SparseClauseContainer() {
+        clear();
+    }
+
+
+
+    [[nodiscard]] int n_classes() const {
+        return classes.size();
+    }
+
+    [[nodiscard]] const std::vector<int>& get_classes() const {
+        return classes;
+    }
+
+    [[nodiscard]] size_t size() const {
+        return classes.size();
+    }
+
+    void populate(const std::vector<int>& classes_to_populate) {
+        for (auto c : classes_to_populate) {
+            this->d.emplace(c, std::make_shared<ClauseType>(*template_instance));
+            this->classes.push_back(c);
+        }
+    }
+
+    std::shared_ptr<ClauseType> operator[](int key) {
+        try{
+            return d.at(key);
+        }catch(const std::out_of_range& e){
+            std::cerr << "Key not found in SparseClauseContainer" << std::endl;
+            std::cerr << e.what() << std::endl;
+            std::cerr << "Key: " << key << std::endl;
+            std::cerr << "Available keys: ";
+            for (auto& item: d) {
+                std::cerr << item.first << " ";
+            }
+            throw e;
+        }
+    }
+
+    ValueIterator<ClauseType> begin() { return ValueIterator<ClauseType>(d.begin()); }
+    ValueIterator<ClauseType> end() { return ValueIterator<ClauseType>(d.end()); }
+
+
+    tl::optional<int> sample(const std::set<uint32_t>& exclude = {}) {
+        // Handle cases where sampling is not possible
+        if (n_classes() <= static_cast<int>(exclude.size()) || n_classes() == 1) {
+            return tl::nullopt; // Indicates no valid sample found
+        }
+
+        while (true) {
+            int idx = pcg32_fast() % n_classes();
+            int sampled_class = classes[idx];
+            if (exclude.find(sampled_class) == exclude.end()) {
+                return sampled_class; // Return immediately once a valid class is found
+            }
+        }
+    }
+
+
+
+
+    void insert(int key, std::shared_ptr<ClauseType> value) {
+        if (d.find(key) == d.end()) {
+            classes.push_back(key);
+        }
+        d[key] = value;
+    }
+
+
+
+
+
+    void clear() {
+        d.clear();
+        classes.clear();
+    }
+};
+
+#endif //TUMLIBPP_SPARSE_CLAUSE_CONTAINER_H
diff --git a/tmu/lib/cpp/include/utils/tm_dataset.h b/tmu/lib/cpp/include/utils/tm_dataset.h
new file mode 100644
index 00000000..1b45b0c5
--- /dev/null
+++ b/tmu/lib/cpp/include/utils/tm_dataset.h
@@ -0,0 +1,98 @@
+//
+// Created by per on 3/3/24.
+//
+
+#ifndef TUMLIBPP_TM_DATASET_H
+#define TUMLIBPP_TM_DATASET_H
+
+#include <fstream>
+#include <sstream>
+#include <vector>
+#include <string>
+#include <cstdint> // for uint32_t
+#include <iostream>
+#include <utility> // for std::pair
+#include <numeric> // for std::accumulate
+#include <random> // for std::random_device, std::mt19937, std::uniform_int_distribution
+
+class TMDataset {
+
+public:
+
+    static std::pair<std::vector<uint32_t>, std::vector<int32_t>> read_dataset_from_txt(const std::string& file_path) {
+        std::ifstream file(file_path);
+        if (!file.is_open()) {
+            throw std::runtime_error("Could not open file " + file_path);
+        }
+
+        std::string line;
+        std::vector<uint32_t> dataset_1d;
+        int num_rows = 0;
+        int num_cols = 0;
+
+        while (std::getline(file, line)) {
+            std::istringstream line_stream(line);
+            std::string value;
+            int current_col_count = 0;
+
+            while (std::getline(line_stream, value, ' ')) {
+                if (!value.empty()) { // Check if value is not empty to handle multiple spaces
+                    uint32_t pixel = std::stoul(value); // Convert string to unsigned long, then implicitly to uint32_t
+                    dataset_1d.push_back(pixel);
+                    ++current_col_count;
+                }
+            }
+
+            if (num_rows == 0) {
+                num_cols = current_col_count; // Set the number of columns based on the first row
+            }
+            ++num_rows;
+        }
+
+        file.close();
+        return {dataset_1d, {num_rows, num_cols}};
+    }
+
+    static void generate_pattern_based_dataset(std::vector<std::vector<uint32_t>>& X, std::vector<uint32_t>& y, int num_samples) {
+        X.clear();
+        y.clear();
+        std::random_device rd;
+        std::mt19937 gen(rd());
+        std::uniform_int_distribution<> distrib(0, 1);
+
+        for (int i = 0; i < num_samples; ++i) {
+            std::vector<uint32_t> x(50);
+            // Randomly generate binary features
+            for (auto& feature : x) {
+                feature = distrib(gen);
+            }
+
+            // Determine the class based on a pattern in the features
+            // Here, we use the sum of '1's in each 10-feature group to influence the class
+            int sum[5] = {0};
+            for (int group = 0; group < 5; ++group) {
+                sum[group] = std::accumulate(x.begin() + group * 10, x.begin() + (group + 1) * 10, 0);
+            }
+
+            // Example pattern-based class determination:
+            // Use the parity of the sum of '1's in each group to determine the class
+            uint32_t class_label = 0;
+            for (int group = 0; group < 5; ++group) {
+                if (sum[group] % 2 == 0) { // If the sum of '1's is even, add to the class label
+                    class_label += 1 << group;
+                }
+            }
+
+            // This scheme can produce 32 different classes (from 0 to 31) based on the combination of even/odd sums.
+            // To limit it to 10 classes, we can simply take the modulo 10 of the class_label.
+            class_label %= 10;
+
+            X.push_back(x);
+            y.push_back(class_label);
+        }
+    }
+
+};
+
+
+#endif //TUMLIBPP_TM_DATASET_H
diff --git a/tmu/lib/cpp/include/utils/tm_math.h b/tmu/lib/cpp/include/utils/tm_math.h
new file mode 100644
index 00000000..7e8efe33
--- /dev/null
+++ b/tmu/lib/cpp/include/utils/tm_math.h
@@ -0,0 +1,55 @@
+//
+// Created by per on 3/3/24.
+//
+
+#ifndef TUMLIBPP_TM_MATH_H
+#define TUMLIBPP_TM_MATH_H
+#include <vector>
+#include <span>
+#include <algorithm>
+
+class TMMath {
+    public:
+
+
+    template<typename T>
+    static T clamp(T value, T low, T high) {
+        return std::max(low, std::min(value, high));
+    }
+
+
+    static void aRange(size_t num_samples, bool shuffle, std::vector<int>& sample_indices) {
+        // Resize the vector to hold all indices
+        sample_indices.resize(num_samples);
+
+        // Fill the vector with consecutive integers starting from 0
+        std::iota(sample_indices.begin(), sample_indices.end(), 0);
+
+        // Shuffle the indices if required
+        if (shuffle) {
+            // Use a random device to seed a random number generator
+            std::random_device rd;
+            std::mt19937 g(rd());
+
+            std::shuffle(sample_indices.begin(), sample_indices.end(), g);
+        }
+    }
+
+
+    template<typename T, typename U>
+    static const std::vector<T> elementwise_multiply(const tcb::span<const T> span1, const tcb::span<const U> span2) {
+        assert(span1.size() == span2.size());
+
+        std::vector<T> result(span1.size());
+
+        for (std::size_t i = 0; i < span1.size(); ++i) {
+            result[i] = span1[i] * span2[i];
+        }
+
+        return result;
+    }
+
+
+};
+
+#endif //TUMLIBPP_TM_MATH_H
diff --git a/tmu/lib/cpp/main.cpp b/tmu/lib/cpp/main.cpp
new file mode 100644
index 00000000..5412ef0d
--- /dev/null
+++ b/tmu/lib/cpp/main.cpp
@@ -0,0 +1,129 @@
+#include "tm_clause_dense.h"
+#include "utils/tm_dataset.h"
+#include <chrono>
+#include <vector>
+#include <tcb/span.hpp>
+#include "models/classifiers/tm_vanilla.h"
+#include <cstdint>
+#include <type_traits>
+static_assert(sizeof(uint32_t) == sizeof(unsigned int), "uint32_t is not the same size as unsigned int");
+static_assert(std::is_unsigned<uint32_t>::value, "uint32_t is not unsigned");
+
+
+void perform_epoch(
+        TMVanillaClassifier<uint32_t>& classifier,
+        int epoch,
+        const tcb::span<uint32_t>& y_train_span,
+        const tcb::span<uint32_t>& x_train_span,
+        const std::vector<int32_t>& x_train_shape,
+        const tcb::span<uint32_t>& x_test_span,
+        const std::vector<int32_t>& x_test_shape,
+        const tcb::span<uint32_t>& y_test_span
+) {
+    auto timer_train_start = std::chrono::high_resolution_clock::now();
+    classifier.fit(
+            y_train_span,
+            x_train_span,
+            x_train_shape,
+            true
+    );
+    auto timer_train_end = std::chrono::high_resolution_clock::now();
+
+    auto timer_pred_start = std::chrono::high_resolution_clock::now();
+    auto predict_result= classifier.predict(x_test_span, x_test_shape);
+    auto y_pred = std::get<0>(predict_result);
+    auto y_pred_shape = std::get<1>(predict_result);
+
+    std::vector<int> results(y_pred.size());
+    std::transform(y_pred.begin(), y_pred.end(), y_test_span.begin(), results.begin(),
+                   [](int pred, int actual) -> int { return pred == actual ? 1 : 0; });
+
+    int correct = std::accumulate(results.begin(), results.end(), 0);
+
+    double accuracy = static_cast<double>(correct) / static_cast<double>(y_test_span.size()) * 100.0;
+    auto timer_pred_end = std::chrono::high_resolution_clock::now();
+
+    std::chrono::duration<double, std::milli> elapsed_train = timer_train_end - timer_train_start;
+    std::chrono::duration<double, std::milli> elapsed_pred = timer_pred_end - timer_pred_start;
+
+    std::cout << "Epoch: " << epoch << " | Accuracy: " << accuracy << "% | "
+              << "Training Time: " << elapsed_train.count() / 1000 << "s | "
+              << "Prediction Time: " << elapsed_pred.count() / 1000 << "s\n";
+}
+
+
+// Function to perform the training and testing of the TMVanillaClassifier
+void train_and_test_classifier() {
+    // Initialize the TMVanillaClassifier with predefined parameters
+    auto NUM_EPOCHS = 10;
+    TMVanillaClassifier<uint32_t> classifier(
+            5000, // T
+            10.0, // s
+            200.0, // d
+            2000, // number_of_clauses
+            false, // _confidence_driven_updating
+            true, // _weighted_clauses
+            true, // type_1
+            true, // type_2
+            false, // type_3
+            1.0,  // _type_i_ii_ratio
+            tl::nullopt, // max_included_literals
+            true,  // boost_true_positive_feedback
+            true, // reuse_random_feedback // TODO - set to true
+            tl::nullopt, // patch_dim
+            8, // number of state bits
+            8, // number of state bits ind
+            100, // batch_size
+            true,
+            42 // _seed
+    );
+
+    // Load datasets
+    auto x_train_result = TMDataset::read_dataset_from_txt("/mnt/disk/git/code/tmu/tmu/lib/cpp/tests/x_train.txt");
+    auto y_train_result = TMDataset::read_dataset_from_txt("/mnt/disk/git/code/tmu/tmu/lib/cpp/tests/y_train.txt");
+    auto x_test_result = TMDataset::read_dataset_from_txt("/mnt/disk/git/code/tmu/tmu/lib/cpp/tests/x_test.txt");
+    auto y_test_result = TMDataset::read_dataset_from_txt("/mnt/disk/git/code/tmu/tmu/lib/cpp/tests/y_test.txt");
+
+    // Assuming read_dataset_from_txt returns a pair or tuple that can be accessed by std::get
+    auto x_train = std::get<0>(x_train_result);
+    auto y_train = std::get<0>(y_train_result);
+    auto x_test = std::get<0>(x_test_result);
+    auto y_test = std::get<0>(y_test_result);
+
+    auto x_train_shape = std::get<1>(x_train_result);
+    auto y_train_shape = std::get<1>(y_train_result);
+    auto x_test_shape = std::get<1>(x_test_result);
+    auto y_test_shape = std::get<1>(y_test_result);
+
+    // Prepare data spans
+    auto x_train_span = tcb::span<uint32_t>(x_train.data(), x_train.size());
+    auto y_train_span = tcb::span<uint32_t>(y_train.data(), y_train.size());
+    auto x_test_span = tcb::span<uint32_t>(x_test.data(), x_test.size());
+    auto y_test_span = tcb::span<uint32_t>(y_test.data(), y_test.size());
+
+
+    // Training and Testing loop
+    for (int epoch = 0; epoch < NUM_EPOCHS; ++epoch) {
+        perform_epoch(
+                classifier,
+                epoch,
+                y_train_span,
+                x_train_span,
+                x_train_shape,
+                x_test_span,
+                x_test_shape,
+                y_test_span
+        );
+    }
+}
+
+
+int main() {
+    try {
+        train_and_test_classifier();
+    } catch (const std::exception& e) {
+        std::cerr << "An error occurred: " << e.what() << std::endl;
+    }
+
+    return 0;
+}
\ No newline at end of file
diff --git a/tmu/lib/cpp/nb/tmulibcpp.cpp b/tmu/lib/cpp/nb/tmulibcpp.cpp
new file mode 100644
index 00000000..fd6aa4a3
--- /dev/null
+++ b/tmu/lib/cpp/nb/tmulibcpp.cpp
@@ -0,0 +1,681 @@
+//
+// Created by per on 3/1/24.
+//
+
+#include <nanobind/nanobind.h>
+#include <nanobind/stl/optional.h>
+#include <nanobind/stl/vector.h>
+#include <nanobind/ndarray.h>
+#include <nanobind/stl/tuple.h>
+#include <nanobind/stl/shared_ptr.h>
+#include <nanobind/stl/set.h>
+#include <nanobind/stl/map.h>
+#include <nanobind/stl/vector.h>
+#include <nanobind/stl/string.h>
+
+#include <iostream>
+
+#include "tm_memory.h"
+#include "tm_clause_dense.h"
+#include "tm_weight_bank.h"
+#include "models/classifiers/tm_vanilla.h"
+#include "utils/sparse_clause_container.h"
+#include <tl/optional.hpp>
+
+
+extern "C" {
+    #include "ClauseBank.h"
+    #include"Tools.h"
+}
+
+using namespace nanobind;
+using namespace nanobind::literals;
+namespace nb = nanobind;
+
+
+NAMESPACE_BEGIN(NB_NAMESPACE)
+    NAMESPACE_BEGIN(detail)
+
+        template <typename T> struct remove_opt_mono<tl::optional<T>>
+                : remove_opt_mono<T> { };
+
+        template <typename T>
+        struct type_caster<tl::optional<T>> {
+            using Caster = make_caster<T>;
+
+            NB_TYPE_CASTER(tl::optional<T>, const_name("Optional[") +
+                                            concat(Caster::Name) +
+                                            const_name("]"))
+
+            type_caster() : value(tl::nullopt) { }
+
+            bool from_python(handle src, uint8_t flags, cleanup_list* cleanup) noexcept {
+                if (src.is_none()) {
+                    value = tl::nullopt;
+                    return true;
+                }
+
+                Caster caster;
+                if (!caster.from_python(src, flags, cleanup))
+                    return false;
+
+                static_assert(
+                        !std::is_pointer_v<T> || is_base_caster_v<Caster>,
+                        "Binding ``optional<T*>`` requires that ``T`` is handled "
+                        "by nanobind's regular class binding mechanism. However, a "
+                        "type caster was registered to intercept this particular "
+                        "type, which is not allowed.");
+
+                value.emplace(caster.operator cast_t<T>());
+
+                return true;
+            }
+
+            template <typename T_>
+            static handle from_cpp(T_ &&value, rv_policy policy, cleanup_list *cleanup) noexcept {
+                if (!value)
+                    return none().release();
+
+                return Caster::from_cpp(forward_like<T_>(*value), policy, cleanup);
+            }
+        };
+
+        template <> struct type_caster<tl::nullopt_t> {
+            bool from_python(handle src, uint8_t, cleanup_list *) noexcept {
+                if (src.is_none())
+                    return true;
+                return false;
+            }
+
+            static handle from_cpp(tl::nullopt_t, rv_policy, cleanup_list *) noexcept {
+                return none().release();
+            }
+
+            NB_TYPE_CASTER(tl::nullopt_t, const_name("None"))
+        };
+
+    NAMESPACE_END(detail)
+NAMESPACE_END(NB_NAMESPACE)
+
+
+template <typename T>
+nb::class_<SparseClauseContainer<T>>& bind_sparse_container(const char *name, nb::module_& m){
+    return nb::class_<SparseClauseContainer<T>>(m, name)
+            .def("__getitem__", &SparseClauseContainer<T>::operator[])
+            .def("__len__", &SparseClauseContainer<T>::size)
+            .def("classes", &SparseClauseContainer<T>::get_classes)
+            .def("sample", &SparseClauseContainer<T>::sample,  "exclude"_a = std::set<int>())
+            .def("set_clause_init", [](
+                    SparseClauseContainer<T>& self,
+                    const nb::object& cls,
+                    const nb::dict& args
+            ) {
+                auto instance = cls(**args); // Call the Python callable to create an instance
+                auto ins = nb::cast<std::shared_ptr<T>>(instance);
+                self.template_instance = ins;
+
+            })
+            .def("populate", &SparseClauseContainer<T>::populate);
+}
+
+
+
+NB_MODULE(tmulibpy, m) {
+
+    bind_sparse_container<TMClauseBankDense<uint32_t>>("SparseClauseContainer", m)
+            .def(nb::init<unsigned int>(), "random_seed"_a)
+
+    ;
+    bind_sparse_container<TMWeightBank<uint32_t>>("SparseWeightContainer", m)
+            .def(nb::init<unsigned int>(), "random_seed"_a)
+    ;
+
+    nanobind::class_<TMMemory<uint32_t>>(m, "TMMemory")
+        .def(nanobind::init<>())
+        .def("reserve", &TMMemory<uint32_t>::reserve)
+        .def("push_back", &TMMemory<uint32_t>::push_back)
+    ;
+
+    nb::class_<TMVanillaClassifier<uint32_t>>(m, "TMVanillaClassifier")
+        .def(nb::init<
+            int,
+            float,
+            float,
+            uint32_t,
+            bool,
+            bool,
+            bool,
+            bool,
+            bool,
+            float,
+            tl::optional<std::size_t>,
+            bool,
+            bool,
+            tl::optional<std::vector<int>>,
+            int32_t ,
+            int32_t,
+            int32_t,
+            bool,
+            int
+        >(),
+            "T"_a,
+            "s"_a,
+            "d"_a,
+            "number_of_clauses"_a,
+            "confidence_driven_updating"_a,
+            "weighted_clauses"_a,
+            "type_i_feedback"_a,
+            "type_ii_feedback"_a,
+            "type_iii_feedback"_a,
+            "type_i_ii_ratio"_a,
+            "max_included_literals"_a,
+            "boost_true_positive_feedback"_a,
+            "reuse_random_feedback"_a,
+            "patch_dim"_a = std::nullopt,
+            "number_of_state_bits"_a,
+            "number_of_state_bits_ind"_a,
+            "batch_size"_a,
+            "incremental"_a,
+            "seed"_a
+        )
+        .def_ro("memory", &TMVanillaClassifier<uint32_t>::memory)
+        .def("get_required_memory_size", &TMVanillaClassifier<uint32_t>::get_required_memory_size)
+        .def("init_after", [](TMVanillaClassifier<uint32_t>& self, const nb::ndarray<>& X, const nb::ndarray<>& Y){
+            std::cout << "init_after (TODO handling numpy arrays)" << std::endl;
+            self.init_after();
+        })
+        .def("initialize", &TMVanillaClassifier<uint32_t>::initialize)
+        .def("init", [](TMVanillaClassifier<uint32_t>& self, nb::ndarray<uint32_t>& X, nb::ndarray<uint32_t>& Y){
+
+            std::vector<int> X_shape = {static_cast<int>(X.shape(0)), static_cast<int>(X.shape(1))};
+            auto x = tcb::span(X.data(), X.size());
+            auto y = tcb::span(Y.data(), Y.size());
+
+            self.init(y, x, X_shape);
+        })
+        .def(
+            "mechanism_compute_update_probabilities",
+            &TMVanillaClassifier<uint32_t>::mechanism_compute_update_probabilities,
+            "is_target"_a,
+            "class_sum"_a
+        )
+        .def(
+        "mechanism_feedback",
+        [](
+                TMVanillaClassifier<uint32_t>& self,
+                bool is_target,
+                uint32_t target,
+                nanobind::ndarray<uint32_t , nb::ndim<1>, c_contig>& clause_outputs,
+                float update_p,
+                nanobind::ndarray<uint32_t, nb::ndim<2>, c_contig>& clause_active,
+                nanobind::ndarray<uint32_t, nb::ndim<1>, c_contig>& literal_active,
+                nanobind::ndarray<uint32_t, nb::ndim<2>, c_contig>& encoded_X_train,
+                int sample_idx
+        ){
+
+            uint32_t* encoded_xi = &encoded_X_train(sample_idx, 0);
+            auto encoded_xi_len = encoded_X_train.shape(1);
+            auto span_encoded_xi = tcb::span(encoded_xi, encoded_xi_len);
+            auto span_clause_outputs = tcb::span(clause_outputs.data(), clause_outputs.size());
+            auto span_clause_active = tcb::span(clause_active.data(), clause_active.size());
+            auto span_literal_active = tcb::span(literal_active.data(), literal_active.size());
+
+            self.mechanism_feedback(
+                    is_target,
+                    target,
+                    span_clause_outputs,
+                    update_p,
+                    span_clause_active,
+                    span_literal_active,
+                    span_encoded_xi
+            );
+
+        },
+        "is_target"_a,
+        "target"_a,
+        "clause_outputs"_a,
+        "update_p"_a,
+        "clause_active"_a,
+        "literal_active"_a,
+        "encoded_X_train"_a,
+        "sample_idx"_a
+        )
+        .def("predict", [](
+                TMVanillaClassifier<uint32_t>& self,
+                nanobind::ndarray<uint32_t, nb::ndim<2>, c_contig>& x_test,
+                bool clip_class_sum = false,
+                bool return_class_sum = false) {
+
+            const auto x_test_span = tcb::span(x_test.data(), x_test.size());
+            const std::vector<int> X_shape = {
+                    static_cast<int>(x_test.shape(0)),
+                    static_cast<int>(x_test.shape(1))
+            };
+
+            auto [argmax, class_sums] = self.predict(
+                    x_test_span,
+                    X_shape,
+                    clip_class_sum,
+                    return_class_sum
+            );
+
+            // TODO return class sums
+            return std::make_tuple(argmax, argmax);
+        },
+        "x"_a,
+        "clip_class_sum"_a = true,
+        "return_class_sum"_a = false,
+         nb::rv_policy::take_ownership)
+
+
+
+            .def("fit",
+        [](
+                TMVanillaClassifier<uint32_t>& self,
+                nanobind::ndarray<uint32_t, nb::ndim<2>, c_contig>& x,
+                nanobind::ndarray<uint32_t, nb::ndim<1>, c_contig>& y,
+                bool shuffle,
+                std::vector<std::string>& metrics
+        ) {
+
+            const auto y_span = tcb::span(y.data(), y.size());
+            const auto x_train_span = tcb::span(x.data(), x.size());
+
+            const std::vector<int> X_shape = {
+                    static_cast<int>(x.shape(0)),
+                    static_cast<int>(x.shape(1))
+            };
+            self.fit(
+                    y_span,
+                    x_train_span,
+                    X_shape,
+                    shuffle
+            );
+
+
+        },
+        "x"_a,
+        "y"_a,
+        "shuffle"_a = true,
+        "metrics"_a,
+         nb::call_guard<nb::gil_scoped_release>()
+        )
+
+        .def("predict_compute_class_sums", [](
+                TMVanillaClassifier<uint32_t>& self,
+                nanobind::ndarray<uint32_t, nb::ndim<2>, c_contig>& encoded_X_test,
+                int ith_sample,
+                bool clip_class_sum
+        ) {
+
+            auto encoded_X_test_span = tcb::span(encoded_X_test.data(), encoded_X_test.size());
+            auto num_items = encoded_X_test.shape(0);
+
+            return self.predict_compute_class_sums(
+                    encoded_X_test_span,
+                    ith_sample,
+                    num_items,
+                    clip_class_sum
+            );
+
+
+        }, "encoded_X_test"_a, "ith_sample"_a, "clip_class_sum"_a)
+
+        .def_prop_ro("positive_clauses", [](TMVanillaClassifier<uint32_t>& self) {
+            return nb::ndarray<nb::numpy, uint32_t>(
+                    self.positive_clauses.data(),
+                    {static_cast<unsigned long>(self.number_of_clauses)}
+            );
+        }, nb::rv_policy::reference)
+        .def_prop_ro("negative_clauses", [](TMVanillaClassifier<uint32_t>& self) {
+            return nb::ndarray<nb::numpy, uint32_t>(
+                    self.negative_clauses.data(),
+                    {static_cast<unsigned long>(self.number_of_clauses)}
+            );
+        }, nb::rv_policy::reference)
+        .def_prop_ro("clause_banks", [](TMVanillaClassifier<uint32_t>& self) {
+            return &self.clause_banks;
+        }, nb::rv_policy::reference)
+        .def_prop_ro("weight_banks", [](TMVanillaClassifier<uint32_t>& self) {
+            return &self.weight_banks;
+        }, nb::rv_policy::reference)
+
+        ;
+
+
+    nb::class_<TMWeightBank<uint32_t>>(m, "TMWeightBank")
+            .def(nb::init<>())
+            .def("initialize", &TMWeightBank<uint32_t>::initialize)
+            .def("get_weights", [](TMWeightBank<uint32_t>& self) {
+                return nb::ndarray<nb::numpy, uint32_t>(
+                        self.weights.data(),
+                        {static_cast<unsigned long>(self.weights.size())}
+                );
+            }, nb::rv_policy::reference)
+            .def("get_required_memory_size", &TMWeightBank<uint32_t>::getRequiredMemorySize)
+            .def("increment", [](
+                    TMWeightBank<uint32_t>& self,
+                    nb::ndarray<nb::numpy, uint32_t, nb::ndim<1>, c_contig>& clause_output,
+                    float update_p,
+                    nb::ndarray<nb::numpy, uint32_t, nb::ndim<1>, c_contig>& clause_active,
+                    bool positive_weights
+            ) {
+                tcb::span<uint32_t> clause_output_span(clause_output.data(), clause_output.size());
+                tcb::span<uint32_t> clause_active_span(clause_active.data(), clause_active.size());
+                self.increment(clause_output_span, update_p, clause_active_span, positive_weights);
+            },
+             "clause_output"_a,
+             "update_p"_a,
+             "clause_active"_a,
+             "positive_weights"_a
+            )
+            .def("decrement", [](
+                    TMWeightBank<uint32_t>& self,
+                    nb::ndarray<nb::numpy, uint32_t, nb::ndim<1>, c_contig>& clause_output,
+                    float update_p,
+                    nb::ndarray<nb::numpy, uint32_t, nb::ndim<1>, c_contig>& clause_active,
+                    bool negative_weights
+            ) {
+                tcb::span<uint32_t> clause_output_span(clause_output.data(), clause_output.size());
+                tcb::span<uint32_t> clause_active_span(clause_active.data(), clause_active.size());
+                self.decrement(clause_output_span, update_p, clause_active_span, negative_weights);
+            },
+             "clause_output"_a,
+             "update_p"_a,
+             "clause_active"_a,
+             "negative_weights"_a)
+
+            ;
+
+
+    nanobind::class_<TMClauseBankDense<uint32_t>>(m, "TMClauseBankDense")
+        .def(nanobind::init<
+                     float,
+                     float,
+                     bool,
+                     bool,
+                     std::vector<int>,
+                     tl::optional<std::vector<int>>,
+                     tl::optional<std::size_t>,
+                     std::size_t,
+                     std::size_t,
+                     std::size_t,
+                     std::size_t,
+                     bool,
+                     int
+             >(),
+             "s"_a,
+             "d"_a,
+             "boost_true_positive_feedback"_a,
+             "reuse_random_feedback"_a,
+             "X_shape"_a,
+             "patch_dim"_a = std::nullopt,
+             "max_included_literals"_a = std::nullopt,
+             "number_of_clauses"_a,
+             "number_of_state_bits"_a,
+             "number_of_state_bits_ind"_a,
+             "batch_size"_a,
+             "incremental"_a,
+             "seed"_a = 0
+        )
+        .def("initialize", &TMClauseBankDense<uint32_t>::initialize)
+        .def("set_ta_state", &TMClauseBankDense<uint32_t>::setTAState)
+        .def("get_ta_state", &TMClauseBankDense<uint32_t>::getTAState)
+        .def_ro("clause_output", &TMClauseBankDense<uint32_t>::clause_output)
+        .def_ro("clause_output_batch", &TMClauseBankDense<uint32_t>::clause_output_batch)
+        .def_ro("clause_and_target", &TMClauseBankDense<uint32_t>::clause_and_target)
+        .def_ro("clause_output_patchwise", &TMClauseBankDense<uint32_t>::clause_output_patchwise)
+        .def_ro("feedback_to_ta", &TMClauseBankDense<uint32_t>::feedback_to_ta)
+        .def_ro("output_one_patches", &TMClauseBankDense<uint32_t>::output_one_patches)
+        .def_ro("literal_clause_count", &TMClauseBankDense<uint32_t>::literal_clause_count)
+        .def_ro("type_ia_feedback_counter", &TMClauseBankDense<uint32_t>::type_ia_feedback_counter)
+        .def_ro("literal_clause_map", &TMClauseBankDense<uint32_t>::literal_clause_map)
+        .def_ro("literal_clause_map_pos", &TMClauseBankDense<uint32_t>::literal_clause_map_pos)
+        .def_ro("false_literals_per_clause", &TMClauseBankDense<uint32_t>::false_literals_per_clause)
+        .def_ro("previous_xi", &TMClauseBankDense<uint32_t>::previous_xi)
+        .def_ro("clause_bank", &TMClauseBankDense<uint32_t>::clause_bank)
+        .def_ro("actions", &TMClauseBankDense<uint32_t>::actions)
+        .def_ro("clause_bank_ind", &TMClauseBankDense<uint32_t>::clause_bank_ind)
+        .def_ro("number_of_ta_chunks", &TMClauseBankDense<uint32_t>::number_of_ta_chunks)
+        .def_ro("number_of_patches", &TMClauseBankDense<uint32_t>::number_of_patches)
+        .def_ro("dim", &TMClauseBankDense<uint32_t>::dim)
+        .def_ro("patch_dim", &TMClauseBankDense<uint32_t>::patch_dim)
+        .def_ro("number_of_clauses", &TMClauseBankDense<uint32_t>::number_of_clauses)
+        .def_ro("number_of_state_bits", &TMClauseBankDense<uint32_t>::number_of_state_bits)
+        .def_ro("number_of_state_bits_ind", &TMClauseBankDense<uint32_t>::number_of_state_bits_ind)
+        .def_ro("batch_size", &TMClauseBankDense<uint32_t>::batch_size)
+        .def_ro("number_of_literals", &TMClauseBankDense<uint32_t>::number_of_literals)
+
+
+        .def("get_clause_bank", [](TMClauseBankDense<uint32_t>& self) {
+            return nb::ndarray<nb::numpy, uint32_t, nb::c_contig>(
+                    self.clause_bank.data(),
+                    {static_cast<unsigned long>(self.calculateClauseBankSize())}
+            );
+        }, nb::rv_policy::reference)
+
+        .def("get_clause_bank_ind", [](TMClauseBankDense<uint32_t>& self) {
+            return nb::ndarray<nb::numpy, uint32_t>(
+                    self.clause_bank_ind.data(),
+                    {static_cast<unsigned long>(self.calculateClauseBankSize())}
+            );
+        }, nb::rv_policy::reference)
+
+        .def("get_clause_output", [](TMClauseBankDense<uint32_t>& self) {
+            return nb::ndarray<nb::numpy, uint32_t>(
+                    self.clause_output.data(),
+                    {static_cast<unsigned long>(self.number_of_clauses)}
+            );
+        }, nb::rv_policy::reference)
+
+        .def("get_clause_output_batch", [](TMClauseBankDense<uint32_t>& self) {
+            return nb::ndarray<nb::numpy, uint32_t>(
+                    self.clause_output_batch.data(),
+                    {static_cast<unsigned long>(self.calculateClauseOutputBatchSize())}
+            );
+        }, nb::rv_policy::reference)
+
+        .def("get_clause_output_patchwise", [](TMClauseBankDense<uint32_t>& self) {
+            return nb::ndarray<nb::numpy, uint32_t>(
+                    self.clause_output_patchwise.data(),
+                    {static_cast<unsigned long>(self.number_of_clauses * self.number_of_patches)}
+            );
+        }, nb::rv_policy::reference)
+
+        .def("get_feedback_to_ta", [](TMClauseBankDense<uint32_t>& self) {
+            return nb::ndarray<nb::numpy, uint32_t>(
+                    self.feedback_to_ta.data(),
+                    {static_cast<unsigned long>(self.number_of_clauses)}
+            );
+        }, nb::rv_policy::reference)
+
+        .def("get_required_memory_size", &TMClauseBankDense<uint32_t>::getRequiredMemorySize)
+        .def("prepare_X", [](
+                     TMClauseBankDense<uint32_t>& self,
+                     nb::ndarray<uint32_t>& x) {
+
+                 if (x.ndim() != 2) { // Assuming x should be a 2D array, adjust as needed
+                     throw std::runtime_error("Input array x must be a 2D array");
+                 }
+
+                 auto number_of_examples = x.shape(0);
+
+                 auto* encoded_X = new uint32_t[number_of_examples * self.number_of_patches * self.number_of_ta_chunks];
+                 nb::capsule owner(encoded_X, [](void* p) noexcept {
+                     delete[] static_cast<uint32_t*>(p);
+                 });
+
+
+                 tmu_encode(
+                         x.data(),
+                         encoded_X,
+                         static_cast<int>(number_of_examples),
+                         std::get<0>(self.dim),
+                         std::get<1>(self.dim),
+                         std::get<2>(self.dim),
+                         std::get<0>(self.patch_dim),
+                         std::get<1>(self.patch_dim),
+                         1,
+                         0
+                 );
+
+                 return nb::ndarray<nb::numpy, uint32_t, c_contig>(
+                         encoded_X,
+                         {number_of_examples, self.number_of_patches * self.number_of_ta_chunks},
+                         owner
+                 );
+             },
+             "x"_a, nb::rv_policy::take_ownership)
+
+
+         .def("calculate_clause_outputs_update", [](
+                TMClauseBankDense<uint32_t>& self,
+                nanobind::ndarray<uint32_t, nb::ndim<1>, c_contig>& literal_active,
+                nanobind::ndarray<uint32_t, nb::ndim<2>, c_contig>& encoded_X,
+                int e){
+
+            auto encoded_xi = &encoded_X(e, 0);
+            auto encoded_xi_len = encoded_X.shape(1);
+            auto span_encoded_xi = tcb::span(encoded_xi, encoded_xi_len);
+
+            auto literal_active_span = tcb::span(literal_active.data(), literal_active.size());
+
+            self.calculate_clause_outputs_update(
+                    literal_active_span,
+                    span_encoded_xi
+            );
+
+
+            return nb::ndarray<nb::numpy, uint32_t, nb::ndim<1>, nb::c_contig>(
+                  self.clause_output.data(),
+                  {static_cast<unsigned long>(self.number_of_clauses)
+            });
+
+         }
+         , "literal_active"_a
+         , "encoded_X"_a
+         , "e"_a
+         , nb::rv_policy::automatic_reference)
+
+         .def("type_i_feedback", [](
+                 TMClauseBankDense<uint32_t>& self,
+                 float update_p,
+                 nb::ndarray<uint32_t>& clause_active,
+                 nb::ndarray<uint32_t>& literal_active,
+                 nanobind::ndarray<uint32_t, nb::ndim<2>, c_contig>& encoded_X,
+                 int e
+                 ) {
+
+            auto clause_active_span = tcb::span(clause_active.data(), clause_active.size());
+            auto literal_active_span = tcb::span(literal_active.data(), literal_active.size());
+            auto encoded_xi = &encoded_X(e, 0);
+            auto encoded_xi_len = encoded_X.shape(1);
+            auto span_encoded_xi = tcb::span(encoded_xi, encoded_xi_len);
+
+            self.type_i_feedback(
+                    update_p,
+                    clause_active_span,
+                    literal_active_span,
+                    span_encoded_xi
+            );
+         }
+            , "update_p"_a
+            , "clause_active"_a
+            , "literal_active"_a
+            , "encoded_X"_a
+            , "e"_a
+         , nb::rv_policy::automatic
+         )
+
+    .def("type_ii_feedback", [](
+                 TMClauseBankDense<uint32_t>& self,
+                 float update_p,
+                 nb::ndarray<uint32_t>& clause_active,
+                 nb::ndarray<uint32_t>& literal_active,
+                 nanobind::ndarray<uint32_t, nb::ndim<2>, nb::c_contig>& encoded_X,
+                 uint32_t target) {
+
+                auto clause_active_span = tcb::span(clause_active.data(), clause_active.size());
+                auto literal_active_span = tcb::span(literal_active.data(), literal_active.size());
+                auto encoded_xi = &encoded_X(target, 0);
+                auto encoded_xi_len = encoded_X.shape(1);
+                auto span_encoded_xi = tcb::span(encoded_xi, encoded_xi_len);
+
+                 self.type_ii_feedback(
+                    update_p,
+                    clause_active_span,
+                    literal_active_span,
+                    span_encoded_xi
+                 );
+         }
+
+            , "update_p"_a
+            , "clause_active"_a
+            , "literal_active"_a
+            , "encoded_X"_a
+            , "e"_a
+            , nb::rv_policy::automatic
+    )
+
+    .def("calculate_clause_outputs_update", [](
+            TMClauseBankDense<uint32_t>& self,
+                nanobind::ndarray<uint32_t>& literal_active,
+                nanobind::ndarray<uint32_t, nb::ndim<2>, c_contig>& encoded_X,
+                uint32_t e) {
+
+        auto encoded_xi = &encoded_X(e, 0);
+        auto encoded_xi_len = encoded_X.shape(1);
+        auto span_encoded_xi = tcb::span(encoded_xi, encoded_xi_len);
+        auto literal_active_span = tcb::span(literal_active.data(), literal_active.size());
+
+        self.calculate_clause_outputs_update(
+                literal_active_span,
+                span_encoded_xi
+        );
+
+         return nb::ndarray<nb::numpy, uint32_t, nb::ndim<1>, nb::c_contig>(
+                 self.clause_output.data(),
+                 {self.number_of_clauses}
+         );
+    }
+    , "literal_active"_a
+    , "encoded_X"_a
+    , "e"_a
+    , nb::rv_policy::automatic_reference)
+
+
+    .def("calculate_clause_outputs_predict", [](
+                TMClauseBankDense<uint32_t>& self,
+                nanobind::ndarray<uint32_t, nb::ndim<2>, c_contig>& encoded_X,
+                int sample_index) {
+
+        auto encoded_xi = &encoded_X(sample_index, 0);
+        auto encoded_xi_len = encoded_X.shape(1);
+        auto n_items = encoded_X.shape(0);
+        auto span_encoded_xi = tcb::span(encoded_xi, encoded_xi_len);
+
+        auto clause_outputs = self.calculate_clause_outputs_predict(
+                span_encoded_xi,
+                sample_index,
+                n_items
+        );
+
+         auto arr = nb::ndarray<nb::numpy, uint32_t , nb::ndim<1>, nb::c_contig>(
+                 clause_outputs.data(),
+                 {self.number_of_clauses}
+         );
+
+        return arr;
+
+    }
+    , "encoded_X"_a
+    , "e"_a
+    , nb::rv_policy::reference)
+
+
+    ;
+
+
+
+
+}
\ No newline at end of file
diff --git a/tmu/lib/cpp/scripts/STM32L475VGTX_FLASH.ld b/tmu/lib/cpp/scripts/STM32L475VGTX_FLASH.ld
new file mode 100644
index 00000000..7e422e19
--- /dev/null
+++ b/tmu/lib/cpp/scripts/STM32L475VGTX_FLASH.ld
@@ -0,0 +1,178 @@
+/**
+ ******************************************************************************
+ * @file      LinkerScript.ld
+ https://raw.githubusercontent.com/efenocchi/STM32-new_project/master/STM32L475VGTX_FLASH.ld
+ * @author    Auto-generated by STM32CubeIDE
+ *  Abstract    : Linker script for B-L475E-IOT01A2 Board embedding STM32L475VGTx Device from stm32l4 series
+ *                      1024Kbytes FLASH
+ *                      96Kbytes RAM
+ *                      32Kbytes RAM2
+ *
+ *            Set heap size, stack size and stack location according
+ *            to application requirements.
+ *
+ *            Set memory bank area and size if external memory is used
+ ******************************************************************************
+ * @attention
+ *
+ * <h2><center>&copy; Copyright (c) 2020 STMicroelectronics.
+ * All rights reserved.</center></h2>
+ *
+ * This software component is licensed by ST under BSD 3-Clause license,
+ * the "License"; You may not use this file except in compliance with the
+ * License. You may obtain a copy of the License at:
+ *                        opensource.org/licenses/BSD-3-Clause
+ *
+ ******************************************************************************
+ */
+
+/* Entry Point */
+ENTRY(Reset_Handler)
+
+/* Highest address of the user mode stack */
+_estack = ORIGIN(RAM) + LENGTH(RAM);	/* end of "RAM" Ram type memory */
+
+_Min_Heap_Size = 0x200 ;	/* required amount of heap  */
+_Min_Stack_Size = 0x400 ;	/* required amount of stack */
+
+/* Memories definition */
+MEMORY
+{
+  RAM    (xrw)    : ORIGIN = 0x20000000,   LENGTH = 96K
+  RAM2    (xrw)    : ORIGIN = 0x10000000,   LENGTH = 32K
+  FLASH    (rx)    : ORIGIN = 0x8000000,   LENGTH = 1024K
+}
+
+/* Sections */
+SECTIONS
+{
+  /* The startup code into "FLASH" Rom type memory */
+  .isr_vector :
+  {
+    . = ALIGN(4);
+    KEEP(*(.isr_vector)) /* Startup code */
+    . = ALIGN(4);
+  } >FLASH
+
+  /* The program code and other data into "FLASH" Rom type memory */
+  .text :
+  {
+    . = ALIGN(4);
+    *(.text)           /* .text sections (code) */
+    *(.text*)          /* .text* sections (code) */
+    *(.glue_7)         /* glue arm to thumb code */
+    *(.glue_7t)        /* glue thumb to arm code */
+    *(.eh_frame)
+
+    KEEP (*(.init))
+    KEEP (*(.fini))
+
+    . = ALIGN(4);
+    _etext = .;        /* define a global symbols at end of code */
+  } >FLASH
+
+  /* Constant data into "FLASH" Rom type memory */
+  .rodata :
+  {
+    . = ALIGN(4);
+    *(.rodata)         /* .rodata sections (constants, strings, etc.) */
+    *(.rodata*)        /* .rodata* sections (constants, strings, etc.) */
+    . = ALIGN(4);
+  } >FLASH
+
+  .ARM.extab   : {
+    . = ALIGN(4);
+    *(.ARM.extab* .gnu.linkonce.armextab.*)
+    . = ALIGN(4);
+  } >FLASH
+
+  .ARM : {
+    . = ALIGN(4);
+    __exidx_start = .;
+    *(.ARM.exidx*)
+    __exidx_end = .;
+    . = ALIGN(4);
+  } >FLASH
+
+  .preinit_array     :
+  {
+    . = ALIGN(4);
+    PROVIDE_HIDDEN (__preinit_array_start = .);
+    KEEP (*(.preinit_array*))
+    PROVIDE_HIDDEN (__preinit_array_end = .);
+    . = ALIGN(4);
+  } >FLASH
+
+  .init_array :
+  {
+    . = ALIGN(4);
+    PROVIDE_HIDDEN (__init_array_start = .);
+    KEEP (*(SORT(.init_array.*)))
+    KEEP (*(.init_array*))
+    PROVIDE_HIDDEN (__init_array_end = .);
+    . = ALIGN(4);
+  } >FLASH
+
+  .fini_array :
+  {
+    . = ALIGN(4);
+    PROVIDE_HIDDEN (__fini_array_start = .);
+    KEEP (*(SORT(.fini_array.*)))
+    KEEP (*(.fini_array*))
+    PROVIDE_HIDDEN (__fini_array_end = .);
+    . = ALIGN(4);
+  } >FLASH
+
+  /* Used by the startup to initialize data */
+  _sidata = LOADADDR(.data);
+
+  /* Initialized data sections into "RAM" Ram type memory */
+  .data :
+  {
+    . = ALIGN(4);
+    _sdata = .;        /* create a global symbol at data start */
+    *(.data)           /* .data sections */
+    *(.data*)          /* .data* sections */
+
+    . = ALIGN(4);
+    _edata = .;        /* define a global symbol at data end */
+
+  } >RAM AT> FLASH
+
+  /* Uninitialized data section into "RAM" Ram type memory */
+  . = ALIGN(4);
+  .bss :
+  {
+    /* This is used by the startup in order to initialize the .bss section */
+    _sbss = .;         /* define a global symbol at bss start */
+    __bss_start__ = _sbss;
+    *(.bss)
+    *(.bss*)
+    *(COMMON)
+
+    . = ALIGN(4);
+    _ebss = .;         /* define a global symbol at bss end */
+    __bss_end__ = _ebss;
+  } >RAM
+
+  /* User_heap_stack section, used to check that there is enough "RAM" Ram  type memory left */
+  ._user_heap_stack :
+  {
+    . = ALIGN(8);
+    PROVIDE ( end = . );
+    PROVIDE ( _end = . );
+    . = . + _Min_Heap_Size;
+    . = . + _Min_Stack_Size;
+    . = ALIGN(8);
+  } >RAM
+
+  /* Remove information from the compiler libraries */
+  /DISCARD/ :
+  {
+    libc.a ( * )
+    libm.a ( * )
+    libgcc.a ( * )
+  }
+
+  .ARM.attributes 0 : { *(.ARM.attributes) }
+}
\ No newline at end of file
diff --git a/tmu/lib/cpp/scripts/build_stm32.sh b/tmu/lib/cpp/scripts/build_stm32.sh
new file mode 100644
index 00000000..0b99d2de
--- /dev/null
+++ b/tmu/lib/cpp/scripts/build_stm32.sh
@@ -0,0 +1,31 @@
+#!/bin/bash
+
+# Get the current directory of the script
+DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
+
+# Pull the Docker image with the STM32 toolchain
+#docker pull lpodkalicki/stm32-toolchain
+
+# Move to the project root directory; adjust this as necessary to point to your project's root
+cd "$DIR/../../"
+
+# Create the build directory if it doesn't exist and move into it
+mkdir -p build-stm32
+
+
+
+# Run a Docker container with the project directory mounted into /workspace
+# Adjust the working directory to /workspace/build-stm32 to match the new build directory
+docker run -v "$(pwd)":/workspace -w /workspace -it lpodkalicki/stm32-toolchain /bin/sh -c "\
+ls -l && \
+cd build-stm32 && rm -rf CMakeFiles && rm -f CMakeCache.txt && cd .. && \
+TOOLCHAIN_FILE=\$(realpath ./cpp/scripts/stm32l4_toolchain.cmake) && \
+cmake -S . -B build-stm32 \
+  -DBUILD_STM32=ON \
+  -DBUILD_PYTHON=OFF \
+  -DBUILD_EXECUTABLE=ON \
+  -DMCU_FAMILY=STM32L4 \
+  -DMCU_MODEL=STM32L475VG \
+  -DCMAKE_TOOLCHAIN_FILE=\$TOOLCHAIN_FILE \
+  && cmake --build build-stm32 \
+"
diff --git a/tmu/lib/cpp/scripts/stm32l4_toolchain.cmake b/tmu/lib/cpp/scripts/stm32l4_toolchain.cmake
new file mode 100644
index 00000000..20544178
--- /dev/null
+++ b/tmu/lib/cpp/scripts/stm32l4_toolchain.cmake
@@ -0,0 +1,40 @@
+# Set the system name to Generic, indicating a bare-metal target
+set(CMAKE_SYSTEM_NAME Generic)
+
+
+# Specify the cross-compiler locations
+set(CMAKE_C_COMPILER arm-none-eabi-gcc)
+set(CMAKE_CXX_COMPILER arm-none-eabi-g++)
+
+# Specify the processor and architecture flags
+# Adjust these flags according to your specific STM32L4 model
+set(CPU_FLAGS "-mcpu=cortex-m4 -mthumb -mfpu=fpv4-sp-d16 -mfloat-abi=hard")
+set(COMMON_FLAGS "${CPU_FLAGS} -ffunction-sections -fdata-sections -MD -Wall -ffreestanding -fpermissive")
+set(COMMON_FLAGS "${COMMON_FLAGS} -specs=nosys.specs")
+
+
+set(CMAKE_C_FLAGS_INIT "${COMMON_FLAGS} -std=gnu11")
+set(CMAKE_CXX_FLAGS_INIT "${COMMON_FLAGS} -std=gnu++14")
+set(CMAKE_ASM_FLAGS_INIT "${CPU_FLAGS}")
+
+# Set linker flags
+# You will need to specify the correct linker script for your MCU model
+# Get path of current fil
+get_filename_component(CURRENT_DIR ${CMAKE_CURRENT_LIST_FILE} DIRECTORY)
+
+# Set the linker script
+set(CMAKE_EXE_LINKER_FLAGS_INIT "-T${CURRENT_DIR}/STM32L475VGTX_FLASH.ld -Wl,--gc-sections")
+
+# Optionally, you can set the build type to Release to optimize the binary size and performance
+set(CMAKE_BUILD_TYPE_INIT Release)
+
+# You may need to specify the path to the CMSIS and STM32L4 HAL library include directories
+#include_directories(
+#        /path/to/CMSIS/Include
+#        /path/to/STM32L4xx_HAL_Driver/Inc
+#)
+
+# And, specify library search paths if you are linking against libraries
+#link_directories(
+#        /path/to/your/libraries
+#)
diff --git a/tmu/lib/cpp/src/memory.cpp b/tmu/lib/cpp/src/memory.cpp
new file mode 100644
index 00000000..91efa6df
--- /dev/null
+++ b/tmu/lib/cpp/src/memory.cpp
@@ -0,0 +1,3 @@
+//
+// Created by per on 3/1/24.
+//
diff --git a/tmu/lib/include/ClauseBankSparse.h b/tmu/lib/include/ClauseBankSparse.h
index 7d365335..e4fd32ca 100644
--- a/tmu/lib/include/ClauseBankSparse.h
+++ b/tmu/lib/include/ClauseBankSparse.h
@@ -44,8 +44,8 @@ void cbs_calculate_clause_outputs_predict_packed_X(
     int number_of_clauses,
     int number_of_literals,
     unsigned int *clause_output_batch,
-    unsigned short *clause_bank_included,
-    unsigned short *clause_bank_included_length
+    unsigned int *clause_bank_included,
+    unsigned int *clause_bank_included_length
 );
 
 void cbs_unpack_clause_output(
@@ -70,8 +70,8 @@ void cbs_calculate_clause_outputs_update(
     int number_of_clauses,
     int number_of_literals,
     unsigned int *clause_output,
-    unsigned short *clause_bank_included,
-    unsigned short *clause_bank_included_length
+    unsigned int *clause_bank_included,
+    unsigned int *clause_bank_included_length
 );
 
 void cbs_calculate_clause_outputs_predict(
@@ -79,8 +79,8 @@ void cbs_calculate_clause_outputs_predict(
     int number_of_clauses,
     int number_of_literals,
     unsigned int *clause_output,
-    unsigned short *clause_bank_included,
-    unsigned short *clause_bank_included_length
+    unsigned int *clause_bank_included,
+    unsigned int *clause_bank_included_length
 );
 
 void cbs_type_i_feedback(
@@ -97,12 +97,12 @@ void cbs_type_i_feedback(
     int number_of_clauses,
     int number_of_literals,
     int number_of_states,
-    unsigned short *clause_bank_included,
-    unsigned short *clause_bank_included_length,
-    unsigned short *clause_bank_excluded,
-    unsigned short *clause_bank_excluded_length,
-    unsigned short *clause_bank_unallocated,
-    unsigned short *clause_bank_unallocated_length
+    unsigned int *clause_bank_included,
+    unsigned int *clause_bank_included_length,
+    unsigned int *clause_bank_excluded,
+    unsigned int *clause_bank_excluded_length,
+    unsigned int *clause_bank_unallocated,
+    unsigned int *clause_bank_unallocated_length
 );
 
 void cbs_type_ii_feedback(
@@ -114,8 +114,8 @@ void cbs_type_ii_feedback(
     int number_of_clauses,
     int number_of_literals,
     int number_of_states,
-    unsigned short *clause_bank_included,
-    unsigned short *clause_bank_included_length,
-    unsigned short *clause_bank_excluded,
-    unsigned short *clause_bank_excluded_length
+    unsigned int *clause_bank_included,
+    unsigned int *clause_bank_included_length,
+    unsigned int *clause_bank_excluded,
+    unsigned int *clause_bank_excluded_length
 );
\ No newline at end of file
diff --git a/tmu/lib/include/fast_rand.h b/tmu/lib/include/fast_rand.h
index d6d5c1bf..6a595f86 100644
--- a/tmu/lib/include/fast_rand.h
+++ b/tmu/lib/include/fast_rand.h
@@ -1,6 +1,7 @@
 #ifndef C_BITWISE_TSETLIN_MACHINE_FAST_RAND_H
 #define C_BITWISE_TSETLIN_MACHINE_FAST_RAND_H
 #include <stdint.h>
+#include <math.h>
 
 #define FAST_RAND_MAX UINT32_MAX
 
diff --git a/tmu/lib/main.c b/tmu/lib/main.c
deleted file mode 100644
index 913f34a5..00000000
--- a/tmu/lib/main.c
+++ /dev/null
@@ -1,11 +0,0 @@
-#include <stdio.h>
-#include "include/Attention.h"
-#include "include/ClauseBank.h"
-#include "include/ClauseBankSparse.h"
-#include "include/Tools.h"
-#include "include/WeightBank.h"
-
-int main() {
-    printf("Hello, World!\n");
-    return 0;
-}
diff --git a/tmu/lib/pyproject.toml b/tmu/lib/pyproject.toml
new file mode 100644
index 00000000..baa64860
--- /dev/null
+++ b/tmu/lib/pyproject.toml
@@ -0,0 +1,30 @@
+[build-system]
+requires = [
+    "setuptools>=42",
+    "wheel",
+    "scikit-build", # For CMake support
+    "cmake",
+    "ninja",
+    "nanobind"
+]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "tmulibpy"
+version = "0.1.0"
+description = "A Python wrapper for tmulibpp library"
+authors = [
+    { name = "Your Name", email = "your.email@example.com" }
+]
+license = {"text" = "MIT"}
+requires-python = ">=3.8"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.8",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+]
+
+[project.urls]
+Homepage = "https://yourprojecthomepage.com"
+Repository = "https://github.com/yourusername/tumlibpp"
diff --git a/tmu/lib/setup.py b/tmu/lib/setup.py
new file mode 100644
index 00000000..aeab0617
--- /dev/null
+++ b/tmu/lib/setup.py
@@ -0,0 +1,27 @@
+from skbuild import setup  # This leverages scikit-build for CMake projects
+import sysconfig
+
+site_packages_path = sysconfig.get_paths()["purelib"]
+python_root = sysconfig.get_paths()["data"]
+
+# get Relative path between python root and site-packages
+relative_path = site_packages_path[len(python_root):]
+
+print(f"site_packages_path: {site_packages_path}")
+print(f"python_root: {python_root}")
+print(f"relative_path: {relative_path}")
+setup(
+    name="tmulibpy",
+    version="0.1.0",
+    description="A Python wrapper for tmulibpp library",
+    author="Your Name",
+    author_email="your.email@example.com",
+    url="https://github.com/yourusername/tumlibpp",
+    license="MIT",
+    cmake_args=[
+        '-DCMAKE_BUILD_TYPE=Release',
+        '-DBUILD_PYTHON=ON',
+        '-DBUILD_EXECUTABLE=OFF'
+    ],
+    cmake_install_dir=f".{relative_path}",
+)
diff --git a/tmu/lib/src/ClauseBank.c b/tmu/lib/src/ClauseBank.c
index e6ff6010..2eed01d2 100644
--- a/tmu/lib/src/ClauseBank.c
+++ b/tmu/lib/src/ClauseBank.c
@@ -66,8 +66,6 @@ static inline void cb_inc(unsigned int *ta_state, unsigned int active, int numbe
 
 	carry = active;
 	for (int b = 0; b < number_of_state_bits; ++b) {
-		if (carry == 0)
-			break;
 
 		carry_next = ta_state[b] & carry; // Sets carry bits (overflow) passing on to next bit
 		ta_state[b] = ta_state[b] ^ carry; // Performs increments with XOR
@@ -82,25 +80,29 @@ static inline void cb_inc(unsigned int *ta_state, unsigned int active, int numbe
 }
 
 // Decrement the states of each of those 32 Tsetlin Automata flagged in the active bit vector.
-static inline void cb_dec(unsigned int *ta_state, unsigned int active, int number_of_state_bits)
-{
+static inline void cb_dec(
+        unsigned int *ta_state,
+        unsigned int active,
+        int number_of_state_bits
+){
 	unsigned int carry, carry_next;
+    unsigned int ta_val;
 
 	carry = active;
 	for (int b = 0; b < number_of_state_bits; ++b) {
-		if (carry == 0)
-			break;
 
-		carry_next = (~ta_state[b]) & carry; // Sets carry bits (overflow) passing on to next bit
-		ta_state[b] = ta_state[b] ^ carry; // Performs increments with XOR
-		carry = carry_next;
+        ta_val = ta_state[b];
+        carry_next = (~ta_val) & carry; // Sets carry bits (overflow) passing on to next bit
+        ta_state[b] = ta_val ^ carry; // Performs increments with XOR
+        carry = carry_next;
+
 	}
 
 	if (carry > 0) {
 		for (int b = 0; b < number_of_state_bits; ++b) {
 			ta_state[b] &= ~carry;
 		}
-	} 
+	}
 }
 
 /* Calculate the output of each clause using the actions of each Tsetline Automaton. */
@@ -300,7 +302,7 @@ void cb_type_i_feedback(
 	}
 	unsigned int number_of_ta_chunks = (number_of_literals-1)/32 + 1;
 
-	if (reuse_random_feedback) {
+	if (reuse_random_feedback && s > 1.0) {
 		cb_initialize_random_streams(feedback_to_ta, number_of_literals, number_of_ta_chunks, s);
 	}
 
@@ -316,7 +318,7 @@ void cb_type_i_feedback(
 
 		cb_calculate_clause_output_feedback(&ta_state[clause_pos], output_one_patches, &clause_output, &clause_patch, number_of_ta_chunks, number_of_state_bits, filter, number_of_patches, literal_active, Xi);
 
-		if (!reuse_random_feedback) {
+		if (!reuse_random_feedback && s > 1.0) {
 			cb_initialize_random_streams(feedback_to_ta, number_of_literals, number_of_ta_chunks, s);
 		}
 
@@ -331,7 +333,11 @@ void cb_type_i_feedback(
 					cb_inc(&ta_state[clause_pos + ta_pos], literal_active[k] & Xi[clause_patch*number_of_ta_chunks + k] & (~feedback_to_ta[k]), number_of_state_bits);
 				}
 
-	 			cb_dec(&ta_state[clause_pos + ta_pos], literal_active[k] & (~Xi[clause_patch*number_of_ta_chunks + k]) & feedback_to_ta[k], number_of_state_bits);
+				if (s > 1.0) {
+		 			cb_dec(&ta_state[clause_pos + ta_pos], literal_active[k] & (~Xi[clause_patch*number_of_ta_chunks + k]) & feedback_to_ta[k], number_of_state_bits);
+		 		} else {
+		 			cb_dec(&ta_state[clause_pos + ta_pos], literal_active[k] & (~Xi[clause_patch*number_of_ta_chunks + k]), number_of_state_bits);
+		 		}
 			}
 		} else {
 			// Type Ib Feedback
@@ -339,7 +345,11 @@ void cb_type_i_feedback(
 			for (int k = 0; k < number_of_ta_chunks; ++k) {
 				unsigned int ta_pos = k*number_of_state_bits;
 
-				cb_dec(&ta_state[clause_pos + ta_pos], literal_active[k] & feedback_to_ta[k], number_of_state_bits);
+				if (s > 1.0) {
+					cb_dec(&ta_state[clause_pos + ta_pos], literal_active[k] & feedback_to_ta[k], number_of_state_bits);
+				} else {
+					cb_dec(&ta_state[clause_pos + ta_pos], literal_active[k], number_of_state_bits);
+				}
 			}
 		}
 	}
diff --git a/tmu/lib/src/ClauseBankSparse.c b/tmu/lib/src/ClauseBankSparse.c
index 58886729..45e2a34f 100644
--- a/tmu/lib/src/ClauseBankSparse.c
+++ b/tmu/lib/src/ClauseBankSparse.c
@@ -119,8 +119,8 @@ void cbs_calculate_clause_outputs_update(
         int number_of_clauses,
         int number_of_literals,
         unsigned int *clause_output,
-        unsigned short *clause_bank_included,
-        unsigned short *clause_bank_included_length
+        unsigned int *clause_bank_included,
+        unsigned int *clause_bank_included_length
 )
 {
     for (int j = 0; j < number_of_clauses; ++j) {
@@ -142,8 +142,8 @@ void cbs_calculate_clause_outputs_predict_packed_X(
         int number_of_clauses,
         int number_of_literals,
         unsigned int *clause_output_batch,
-        unsigned short *clause_bank_included,
-        unsigned short *clause_bank_included_length
+        unsigned int *clause_bank_included,
+        unsigned int *clause_bank_included_length
 )
 {
      for (int j = 0; j < number_of_clauses; ++j) {
@@ -165,8 +165,8 @@ void cbs_calculate_clause_outputs_predict(
         int number_of_clauses,
         int number_of_literals,
         unsigned int *clause_output,
-        unsigned short *clause_bank_included,
-        unsigned short *clause_bank_included_length
+        unsigned int *clause_bank_included,
+        unsigned int *clause_bank_included_length
 )
 {
     for (int j = 0; j < number_of_clauses; ++j) {
@@ -202,12 +202,12 @@ void cbs_type_i_feedback(
         int number_of_clauses,
         int number_of_literals,
         int number_of_states,
-        unsigned short *clause_bank_included,
-        unsigned short *clause_bank_included_length,
-        unsigned short *clause_bank_excluded,
-        unsigned short *clause_bank_excluded_length,
-        unsigned short *clause_bank_unallocated,
-        unsigned short *clause_bank_unallocated_length
+        unsigned int *clause_bank_included,
+        unsigned int *clause_bank_included_length,
+        unsigned int *clause_bank_excluded,
+        unsigned int *clause_bank_excluded_length,
+        unsigned int *clause_bank_unallocated,
+        unsigned int *clause_bank_unallocated_length
 )
 {
 	unsigned int number_of_ta_chunks = (number_of_literals-1)/32 + 1;
@@ -307,7 +307,39 @@ void cbs_type_i_feedback(
                 }
             }
         } else {
-        	int k = clause_bank_included_length[j];
+            int k;
+            if (((float)fast_rand())/((float)FAST_RAND_MAX) <= 1.0/feedback_rate_excluded_literals) {
+                k = clause_bank_excluded_length[j];
+                while (k--) {
+                    int clause_excluded_pos = clause_pos_base + k*2;
+                    unsigned int literal_chunk = clause_bank_excluded[clause_excluded_pos] / 32;
+                    unsigned int literal_pos = clause_bank_excluded[clause_excluded_pos] % 32;
+            
+                    if ((literal_active[literal_chunk] & (1U << literal_pos)) == 0) {
+                        continue;
+                    }
+
+                    if ((clause_bank_excluded[clause_excluded_pos + 1] > feedback_rate_excluded_literals) && (((float)fast_rand())/((float)FAST_RAND_MAX) <= 1.0/s)) {
+                        clause_bank_excluded[clause_excluded_pos + 1] -= feedback_rate_excluded_literals;
+
+                        if ((int)clause_bank_excluded[clause_excluded_pos + 1] <= absorbing) {
+                           if (clause_bank_unallocated_length[j] == 0 || literal_insertion_state == -1) {
+                                clause_bank_excluded_length[j] -= 1;
+                                int clause_excluded_end_pos = clause_pos_base + clause_bank_excluded_length[j]*2;
+                                clause_bank_excluded[clause_excluded_pos] = clause_bank_excluded[clause_excluded_end_pos];
+                                clause_bank_excluded[clause_excluded_pos + 1] = clause_bank_excluded[clause_excluded_end_pos + 1];
+                            } else {
+                                clause_bank_unallocated_length[j] -= 1;
+                                int clause_unallocated_end_pos = j*number_of_literals + clause_bank_unallocated_length[j];
+                                clause_bank_excluded[clause_excluded_pos] = clause_bank_unallocated[clause_unallocated_end_pos];
+                                clause_bank_excluded[clause_excluded_pos + 1] = literal_insertion_state;
+                            }
+                        }
+                    }
+                }
+            }
+
+        	k = clause_bank_included_length[j];
 			while (k--) {
 				int clause_included_pos = clause_pos_base + k*2;
             	unsigned int literal_chunk = clause_bank_included[clause_included_pos] / 32;
@@ -332,37 +364,6 @@ void cbs_type_i_feedback(
                     }
                 }
             }
- 			
-            if (((float)fast_rand())/((float)FAST_RAND_MAX) <= 1.0/feedback_rate_excluded_literals) {
-     			k = clause_bank_excluded_length[j];
-    			while (k--) {
-    				int clause_excluded_pos = clause_pos_base + k*2;
-                	unsigned int literal_chunk = clause_bank_excluded[clause_excluded_pos] / 32;
-                	unsigned int literal_pos = clause_bank_excluded[clause_excluded_pos] % 32;
-    		
-                    if ((literal_active[literal_chunk] & (1U << literal_pos)) == 0) {
-                        continue;
-                    }
-
-                	if ((clause_bank_excluded[clause_excluded_pos + 1] > feedback_rate_excluded_literals) && (((float)fast_rand())/((float)FAST_RAND_MAX) <= 1.0/s)) {
-                        clause_bank_excluded[clause_excluded_pos + 1] -= feedback_rate_excluded_literals;
-
-                        if ((int)clause_bank_excluded[clause_excluded_pos + 1] <= absorbing) {
-                           if (clause_bank_unallocated_length[j] == 0 || literal_insertion_state == -1) {
-                                clause_bank_excluded_length[j] -= 1;
-                                int clause_excluded_end_pos = clause_pos_base + clause_bank_excluded_length[j]*2;
-                                clause_bank_excluded[clause_excluded_pos] = clause_bank_excluded[clause_excluded_end_pos];
-                                clause_bank_excluded[clause_excluded_pos + 1] = clause_bank_excluded[clause_excluded_end_pos + 1];
-                            } else {
-                                clause_bank_unallocated_length[j] -= 1;
-                                int clause_unallocated_end_pos = j*number_of_literals + clause_bank_unallocated_length[j];
-                                clause_bank_excluded[clause_excluded_pos] = clause_bank_unallocated[clause_unallocated_end_pos];
-                                clause_bank_excluded[clause_excluded_pos + 1] = literal_insertion_state;
-                            }
-                        }
-                    }
-                }
-            }
     	}
     }
 }
@@ -376,10 +377,10 @@ void cbs_type_ii_feedback(
         int number_of_clauses,
         int number_of_literals,
         int number_of_states,
-        unsigned short *clause_bank_included,
-        unsigned short *clause_bank_included_length,
-        unsigned short *clause_bank_excluded,
-        unsigned short *clause_bank_excluded_length
+        unsigned int *clause_bank_included,
+        unsigned int *clause_bank_included_length,
+        unsigned int *clause_bank_excluded,
+        unsigned int *clause_bank_excluded_length
 )
 {
     for (int j = 0; j < number_of_clauses; ++j) {
diff --git a/tmu/lib/src/Tools.c b/tmu/lib/src/Tools.c
index a5e5a3e6..9514bb2e 100644
--- a/tmu/lib/src/Tools.c
+++ b/tmu/lib/src/Tools.c
@@ -25,20 +25,23 @@ This code implements the Convolutional Tsetlin Machine from paper arXiv:1905.096
 
 */
 
-#include <stdio.h>
 #include <string.h>
 #include <stdlib.h>
 
-unsigned int compareints(const void * a, const void * b)
-{
-  return(*(unsigned int*)a - *(unsigned int*)b);
+int compareints(const void *a, const void *b) {
+    const unsigned int *ia = (const unsigned int *)a;
+    const unsigned int *ib = (const unsigned int *)b;
+    if (*ia < *ib) return -1;
+    if (*ia > *ib) return 1;
+    return 0;
 }
 
+
 void tmu_produce_autoencoder_example(
-        unsigned int *active_output,
+        const unsigned int *active_output,
         int number_of_active_outputs,
-        unsigned int *indptr_row,
-        unsigned int *indices_row,
+        const unsigned int *indptr_row,
+        const unsigned int *indices_row,
         int number_of_rows,
         unsigned int *indptr_col,
         unsigned int *indices_col,
diff --git a/tmu/models/base.py b/tmu/models/base.py
index 675e008b..798bbbae 100644
--- a/tmu/models/base.py
+++ b/tmu/models/base.py
@@ -109,7 +109,7 @@ def __init__(
         self.s = s
 
         self.confidence_driven_updating = confidence_driven_updating
-
+        self.type_i_ii_ratio = type_i_ii_ratio
         if type_i_ii_ratio >= 1.0:
             self.type_i_p = 1.0
             self.type_ii_p = 1.0 / type_i_ii_ratio
@@ -211,8 +211,8 @@ def init(self, X: np.ndarray, Y: np.ndarray):
         if self.initialized:
             return
 
-        self.init_before(X, Y)
         self.number_of_classes = self.init_num_classes(X, Y)
+        self.init_before(X, Y)
         self.init_clause_bank(X, Y)
         self.init_weight_bank(X, Y)
         self.init_after(X, Y)
@@ -288,7 +288,6 @@ def _build_cpu_sparse_bank(self, X: np.ndarray):
     def build_clause_bank(self, X: np.ndarray):
         if self.platform == "CPU":
             clause_bank_type, clause_bank_args = self._build_cpu_bank(X=X)
-
         elif self.platform in ["GPU", "CUDA"]:
             clause_bank_type, clause_bank_args = self._build_gpu_bank(X=X)
         elif self.platform == "CPU_sparse":
diff --git a/tmu/models/classification/vanilla_classifier.py b/tmu/models/classification/vanilla_classifier.py
index 5e2a6c9f..5fc5fa4e 100644
--- a/tmu/models/classification/vanilla_classifier.py
+++ b/tmu/models/classification/vanilla_classifier.py
@@ -16,7 +16,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 import typing
-from collections import defaultdict
+
 
 from tmu.models.base import MultiClauseBankMixin, MultiWeightBankMixin, TMBaseModel
 from tmu.util.encoded_data_cache import DataEncoderCache
@@ -136,7 +136,7 @@ def mechanism_feedback(
             is_target: bool,
             target: int,
             clause_outputs: np.ndarray,
-            update_p: np.ndarray,
+            update_p: float,
             clause_active: np.ndarray,
             literal_active: np.ndarray,
             encoded_X_train: np.ndarray,
@@ -147,6 +147,7 @@ def mechanism_feedback(
 
         if self.weighted_clauses:
             if is_target:
+
                 self.weight_banks[target].increment(
                     clause_output=clause_outputs,
                     update_p=update_p,
@@ -205,7 +206,7 @@ def mechanism_clause_sum(
             literal_active: np.ndarray,
             encoded_X_train: np.ndarray,
             sample_idx: int
-    ) -> typing.Tuple[np.ndarray, np.ndarray]:
+    ) -> typing.Tuple[int, np.ndarray]:
 
         clause_outputs: np.ndarray = self.clause_banks[target].calculate_clause_outputs_update(
             literal_active=literal_active,
@@ -218,14 +219,14 @@ def mechanism_clause_sum(
             clause_outputs
         ).astype(np.int32)
 
-        class_sum: np.ndarray = np.clip(class_sum, -self.T, self.T)
+        class_sum: int = np.clip(class_sum, -self.T, self.T).item()
 
         return class_sum, clause_outputs
 
     def mechanism_compute_update_probabilities(
             self,
             is_target: bool,
-            class_sum: np.ndarray
+            class_sum: int
     ) -> float:
         # Confidence-driven updating method
         if self.confidence_driven_updating:
@@ -283,11 +284,13 @@ def _fit_sample_target(
         :return: Computed update probability for the class.
         """
 
+
         update_p: float = self.mechanism_compute_update_probabilities(
             is_target=is_target_class,
             class_sum=class_sum
         )
 
+
         self.mechanism_feedback(
             is_target=is_target_class,
             target=class_value,
@@ -304,7 +307,7 @@ def _fit_sample_target(
     def _fit_sample(
             self,
             target: int,
-            not_target: int,
+            not_target: int | None,
             sample_idx: int,
             clause_active: np.ndarray,
             literal_active: np.ndarray,
@@ -331,7 +334,7 @@ def _fit_sample(
         )
 
         # for incremental, and when we only have 1 sample, there is no other targets
-        if self.weight_banks.n_classes == 1:
+        if not_target is None:
             return dict(
                 update_p_target=update_p_target,
                 update_p_not_target=None
@@ -363,28 +366,29 @@ def _fit_sample(
 
     def fit(
             self,
-            X: np.ndarray,
-            Y: np.ndarray,
+            X: np.ndarray[np.uint32],
+            Y: np.ndarray[np.uint32],
             shuffle: bool = True,
             metrics: typing.Optional[list] = None,
             *args,
             **kwargs
     ):
         metrics = metrics or []
-        assert len(X) == len(Y), "X and Y must have the same length"
+        assert X.shape[0] == len(Y), "X and Y must have the same number of samples"
         assert len(X.shape) >= 2, "X must be a 2D array"
         assert len(Y.shape) == 1, "Y must be a 1D array"
+        assert X.dtype == np.uint32, "X must be of type uint32"
+        assert Y.dtype == np.uint32, "Y must be of type uint32"
 
         self.init(X, Y)
         self.metrics.clear()
 
-        Ym = Y.astype(np.uint32)
-
         encoded_X_train: np.ndarray = self.train_encoder_cache.get_encoded_data(
             data=X,
             encoder_func=lambda x: self.clause_banks[0].prepare_X(x)
         )
 
+
         clause_active: np.ndarray = self.mechanism_clause_active()
         literal_active: np.ndarray = self.mechanism_literal_active()
 
@@ -393,8 +397,8 @@ def fit(
             self.rng.shuffle(sample_indices)
 
         for sample_idx in sample_indices:
-            target: int = Ym[sample_idx]
-            not_target: int = self.weight_banks.sample(exclude=[target])
+            target: int = Y[sample_idx]
+            not_target: int | None = self.weight_banks.sample(exclude=[target])
 
             history: dict = self._fit_sample(
                 target=target,
diff --git a/tmu/tools.py b/tmu/tools.py
index 4bfbab44..c27c12bd 100644
--- a/tmu/tools.py
+++ b/tmu/tools.py
@@ -90,12 +90,32 @@ def data(self):
     def slice_ptr(self, val):
         return ffi.cast(self._cast_to, self._data[val].ctypes.data)
 
-def encode(X, number_of_examples, number_of_patches, number_of_ta_chunks, dim, patch_dim, class_features,
-           append_negated=True):
+
+def encode(
+        X,
+        number_of_examples,
+        number_of_patches,
+        number_of_ta_chunks,
+        dim,
+        patch_dim,
+        class_features,
+        append_negated=True
+):
+
     Xm = np.ascontiguousarray(X.flatten()).astype(np.uint32)
     encoded_X = np.ascontiguousarray(
         np.empty(int(number_of_examples * number_of_patches * number_of_ta_chunks), dtype=np.uint32))
-    lib.tmu_encode(ffi.cast("unsigned int *", Xm.ctypes.data), ffi.cast("unsigned int *", encoded_X.ctypes.data),
-                   number_of_examples, dim[0], dim[1], dim[2], patch_dim[0], patch_dim[1], int(append_negated),
-                   class_features)
+
+    lib.tmu_encode(
+        ffi.cast("unsigned int *", Xm.ctypes.data),
+        ffi.cast("unsigned int *", encoded_X.ctypes.data),
+        number_of_examples,
+        dim[0],
+        dim[1],
+        dim[2],
+        patch_dim[0],
+        patch_dim[1],
+        int(append_negated),
+        class_features
+    )
     return np.ascontiguousarray(encoded_X.reshape((int(number_of_examples), number_of_patches * number_of_ta_chunks)))
diff --git a/tmu/util/encoded_data_cache.py b/tmu/util/encoded_data_cache.py
index eb84af27..543b5fae 100644
--- a/tmu/util/encoded_data_cache.py
+++ b/tmu/util/encoded_data_cache.py
@@ -1,7 +1,8 @@
-import hashlib
 from typing import Optional
-
+import xxhash
 import numpy as np
+from scipy.sparse import issparse
+
 
 class DataEncoderCache:
     def __init__(self, seed: int):
@@ -10,12 +11,28 @@ def __init__(self, seed: int):
         self.array_hash: Optional[str] = None
         self.encoded_data: Optional[np.ndarray] = None
 
-    def compute_hash(self, arr: np.ndarray) -> str:
-        """Compute a hash for a numpy array."""
+    def compute_hash_csr_matrix(self, csr_mat):
+        # Convert the components of the csr_matrix to bytes
+        data_bytes = csr_mat.data.tobytes()
+        indices_bytes = csr_mat.indices.tobytes()
+        indptr_bytes = csr_mat.indptr.tobytes()
+
+        # Concatenate the bytes representations
+        total_bytes = data_bytes + indices_bytes + indptr_bytes
+
+        # Compute the hash on the concatenated bytes
+        hash_value = xxhash.xxh3_64_hexdigest(total_bytes)
+
+        return hash_value
 
-        sampled_indices = self.rng.choice(arr.size, min(15, arr.size), replace=False)
-        sampled_values = arr.flat[sampled_indices]
-        return hashlib.sha256(sampled_values.tobytes()).hexdigest()
+    def compute_hash(self, arr):
+        """Compute a hash for a numpy array or csr_matrix."""
+        if issparse(arr):
+            # It's a sparse matrix, handle specially
+            return self.compute_hash_csr_matrix(arr)
+        else:
+            # It's a dense array, proceed as before
+            return xxhash.xxh3_64_hexdigest(arr.tobytes())
 
     def get_encoded_data(self, data: np.ndarray, encoder_func) -> np.ndarray:
         """Get encoded data for an array, using cache if available."""
diff --git a/tmu/util/sparse_clause_container.py b/tmu/util/sparse_clause_container.py
index 910b077b..c12a3a08 100644
--- a/tmu/util/sparse_clause_container.py
+++ b/tmu/util/sparse_clause_container.py
@@ -28,7 +28,11 @@ def classes(self):
     def __len__(self):
         return len(self._classes)
 
-    def sample(self, n=1, exclude=None):
+
+    def sample(self, n=1, exclude=None) -> int | None:
+        if self.n_classes == 1:
+            return None
+
         if exclude is None:
             exclude = set()
         else:


Test patch:
diff --git a/.github/workflows/build-tests.yml b/.github/workflows/build-tests.yml
index 9f9c275d..545f8e75 100644
--- a/.github/workflows/build-tests.yml
+++ b/.github/workflows/build-tests.yml
@@ -27,11 +27,12 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install pytest
-          pip install -r examples/requirements.txt
       - name: Install tmu
         run: |
-          pip install develop .
+          pip install .
+          pip install .[composite]
+          pip install .[examples]
+          pip install .[tests]
       - name: Test with pytest
         run: pytest test --doctest-modules --junitxml=junit/test-results-${{ matrix.os }}-${{ matrix.python-version }}.xml
       - name: Upload pytest test results
diff --git a/test/test_classifiers.py b/test/test_classifiers.py
index 1220eef6..82dd1545 100644
--- a/test/test_classifiers.py
+++ b/test/test_classifiers.py
@@ -55,7 +55,7 @@ def setUp(self) -> None:
 class TMMultiChannelClassifierTests(unittest.TestCase, ClassifierTests):
 
     def setUp(self) -> None:
-        from tmu.models.classification.multichannel_classifier import TMMultiChannelClassifier
+        from tmu.experimental.models.multichannel_classifier import TMMultiChannelClassifier
         import numpy as np
         self.model = TMMultiChannelClassifier(
             number_of_clauses=4,
@@ -75,7 +75,7 @@ class TMCoalescedClassifierTests(unittest.TestCase, ClassifierTests):
     def setUp(self) -> None:
         from tmu.models.classification.coalesced_classifier import TMCoalescedClassifier
         self.model = TMCoalescedClassifier(
-            number_of_clauses=4,
+            number_of_clauses=10,
             T=10,
             s=10.0,
             max_included_literals=32,
@@ -88,8 +88,7 @@ def setUp(self) -> None:
 class OneVSOneClassifierTests(unittest.TestCase, ClassifierTests):
 
     def setUp(self) -> None:
-        from tmu.models.classification.one_vs_one_classifier import TMOneVsOneClassifier
-        import numpy as np
+        from tmu.experimental.models.one_vs_one_classifier import TMOneVsOneClassifier
 
         self.model = TMOneVsOneClassifier(
             number_of_clauses=4,
diff --git a/test/test_components.py b/test/test_components.py
new file mode 100644
index 00000000..bf0efd9b
--- /dev/null
+++ b/test/test_components.py
@@ -0,0 +1,80 @@
+import unittest
+from tmu.composite.components import image as image_components
+from tmu.composite.components.base import TMComponent
+from tmu.composite.composite import TMComposite
+from tmu.composite.config import TMClassifierConfig
+from tmu.data.cifar10 import CIFAR10
+from tmu.models.classification.vanilla_classifier import TMClassifier
+
+
+class TestTMCompositeImageComponents(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        # Load CIFAR-10 dataset
+        cifar_ds = CIFAR10()
+        data = cifar_ds.get()
+
+        X_train_org = data["x_train"]
+        Y_train = data["y_train"]
+        X_test_org = data["x_test"]
+        Y_test = data["y_test"]
+
+
+        Y_test = Y_test.reshape(Y_test.shape[0])
+        Y_train = Y_train.reshape(Y_train.shape[0])
+        X_train_org = X_train_org[:50]
+        Y_train = Y_train[:50]
+        X_test_org = X_test_org[:50]
+        Y_test = Y_test[:50]
+
+        # Initialize datasets
+        cls.data_train = dict(X=X_train_org, Y=Y_train)
+        cls.data_test = dict(X=X_test_org, Y=Y_test)
+
+        cls.platform = 'CPU'
+        cls.epochs = 10
+
+    def test_all_components(self):
+        # Dynamically get all classes from tmcomposite.components.image
+        components_classes = [getattr(image_components, attr) for attr in dir(image_components)
+                              if isinstance(getattr(image_components, attr), type) and
+                              issubclass(getattr(image_components, attr), TMComponent) and
+                              getattr(image_components, attr) != TMComponent]
+
+        # Create a composite model using all the loaded classes as components
+        components = []
+        for component_class in components_classes:
+            component = component_class(
+                TMClassifier,
+                TMClassifierConfig(
+                    number_of_clauses=50,
+                    T=20,
+                    s=10.0,
+                    max_included_literals=32,
+                    platform=self.platform,
+                    weighted_clauses=True,
+                    patch_dim=(10, 10)
+                ),
+                epochs=self.epochs
+            )
+            components.append(component)
+
+        composite_model = TMComposite(components=components, use_multiprocessing=False)
+
+        # Train and test
+        composite_model.fit(
+            data=self.data_train,
+            callbacks=[
+            ]
+        )
+
+        preds = composite_model.predict(data=self.data_test)
+        y_true = self.data_test["Y"].flatten()
+
+        for k, v in preds.items():
+            accuracy = (v == y_true).mean()
+            print(f"{k} Accuracy: %.1f" % (100 * accuracy))
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/test/test_datasets.py b/test/test_datasets.py
index 8aedaf94..3ed13f02 100644
--- a/test/test_datasets.py
+++ b/test/test_datasets.py
@@ -1,110 +1,116 @@
-import pytest
+import unittest
 
+class TestMNISTDataset(unittest.TestCase):
 
-class TestMNISTDatasetobject:
-
-    def setup_class(self):
+    @classmethod
+    def setUpClass(cls):
         from tmu.data import MNIST
-        self.dataset_instance = MNIST()
+        cls.dataset_instance = MNIST()
 
-    def teardown_class(self):
-        self.dataset_instance = None
+    @classmethod
+    def tearDownClass(cls):
+        cls.dataset_instance = None
 
     def test_mnist_dict(self):
         dataset = self.dataset_instance.get()
-        assert len(dataset) == 4
-        assert "x_train" in dataset
-        assert "y_train" in dataset
-        assert "x_test" in dataset
-        assert "y_test" in dataset
-        assert len(dataset["x_train"]) == 60000
-        assert len(dataset["y_train"]) == 60000
-        assert len(dataset["x_test"]) == 10000
-        assert len(dataset["y_test"]) == 10000
+        self.assertEqual(len(dataset), 4)
+        self.assertIn("x_train", dataset)
+        self.assertIn("y_train", dataset)
+        self.assertIn("x_test", dataset)
+        self.assertIn("y_test", dataset)
+        self.assertEqual(len(dataset["x_train"]), 60000)
+        self.assertEqual(len(dataset["y_train"]), 60000)
+        self.assertEqual(len(dataset["x_test"]), 10000)
+        self.assertEqual(len(dataset["y_test"]), 10000)
 
     def test_mnist_list(self):
         dataset = self.dataset_instance.get_list()
-        assert len(dataset) == 4
-        assert isinstance(dataset, list)
+        self.assertEqual(len(dataset), 4)
+        self.assertIsInstance(dataset, list)
 
 
-class TestFashionMNISTDataset:
+class TestFashionMNISTDataset(unittest.TestCase):
 
-    def setup_class(self):
+    @classmethod
+    def setUpClass(cls):
         from tmu.data.fashion_mnist import FashionMNIST
-        self.dataset_instance = FashionMNIST()
+        cls.dataset_instance = FashionMNIST()
 
-    def teardown_class(self):
-        self.dataset_instance = None
+    @classmethod
+    def tearDownClass(cls):
+        cls.dataset_instance = None
 
     def test_fashion_mnist_dict(self):
         dataset = self.dataset_instance.get()
-        assert len(dataset) == 4
-        assert "x_train" in dataset
-        assert "y_train" in dataset
-        assert "x_test" in dataset
-        assert "y_test" in dataset
-        assert len(dataset["x_train"]) == 60000
-        assert len(dataset["y_train"]) == 60000
-        assert len(dataset["x_test"]) == 10000
-        assert len(dataset["y_test"]) == 10000
+        self.assertEqual(len(dataset), 4)
+        self.assertIn("x_train", dataset)
+        self.assertIn("y_train", dataset)
+        self.assertIn("x_test", dataset)
+        self.assertIn("y_test", dataset)
+        self.assertEqual(len(dataset["x_train"]), 60000)
+        self.assertEqual(len(dataset["y_train"]), 60000)
+        self.assertEqual(len(dataset["x_test"]), 10000)
+        self.assertEqual(len(dataset["y_test"]), 10000)
 
     def test_fashion_mnist_list(self):
         dataset = self.dataset_instance.get_list()
-        assert len(dataset) == 4
-        assert isinstance(dataset, list)
+        self.assertEqual(len(dataset), 4)
+        self.assertIsInstance(dataset, list)
 
 
-class TestKuzushijiMNISTDataset:
+class TestKuzushijiMNISTDataset(unittest.TestCase):
 
-    def setup_class(self):
+    @classmethod
+    def setUpClass(cls):
         from tmu.data.fashion_mnist import KuzushijiMNIST
-        self.dataset_instance = KuzushijiMNIST()
+        cls.dataset_instance = KuzushijiMNIST()
 
-    def teardown_class(self):
-        self.dataset_instance = None
+    @classmethod
+    def tearDownClass(cls):
+        cls.dataset_instance = None
 
     def test_kuzushiji_mnist_dict(self):
         dataset = self.dataset_instance.get()
-        assert len(dataset) == 4
-        assert "x_train" in dataset
-        assert "y_train" in dataset
-        assert "x_test" in dataset
-        assert "y_test" in dataset
-        assert len(dataset["x_train"]) == 60000
-        assert len(dataset["y_train"]) == 60000
-        assert len(dataset["x_test"]) == 10000
-        assert len(dataset["y_test"]) == 10000
+        self.assertEqual(len(dataset), 4)
+        self.assertIn("x_train", dataset)
+        self.assertIn("y_train", dataset)
+        self.assertIn("x_test", dataset)
+        self.assertIn("y_test", dataset)
+        self.assertEqual(len(dataset["x_train"]), 60000)
+        self.assertEqual(len(dataset["y_train"]), 60000)
+        self.assertEqual(len(dataset["x_test"]), 10000)
+        self.assertEqual(len(dataset["y_test"]), 10000)
 
     def test_kuzushiji_mnist_list(self):
         dataset = self.dataset_instance.get_list()
-        assert len(dataset) == 4
-        assert isinstance(dataset, list)
+        self.assertEqual(len(dataset), 4)
+        self.assertIsInstance(dataset, list)
 
 
-class TestCIFAR100Dataset:
+class TestCIFAR100Dataset(unittest.TestCase):
 
-    def setup_class(self):
+    @classmethod
+    def setUpClass(cls):
         from tmu.data.fashion_mnist import CIFAR100
-        self.dataset_instance = CIFAR100()
+        cls.dataset_instance = CIFAR100()
 
-    def teardown_class(self):
-        self.dataset_instance = None
+    @classmethod
+    def tearDownClass(cls):
+        cls.dataset_instance = None
 
     def test_cifar100_dict(self):
         dataset = self.dataset_instance.get()
-        assert len(dataset) == 4
-        assert "x_train" in dataset
-        assert "y_train" in dataset
-        assert "x_test" in dataset
-        assert "y_test" in dataset
-        # adjust the following counts according to your train-test split
-        assert len(dataset["x_train"]) == 50000
-        assert len(dataset["y_train"]) == 50000
-        assert len(dataset["x_test"]) == 10000
-        assert len(dataset["y_test"]) == 10000
+        self.assertEqual(len(dataset), 4)
+        self.assertIn("x_train", dataset)
+        self.assertIn("y_train", dataset)
+        self.assertIn("x_test", dataset)
+        self.assertIn("y_test", dataset)
+        self.assertEqual(len(dataset["x_train"]), 50000)
+        self.assertEqual(len(dataset["y_train"]), 50000)
+        self.assertEqual(len(dataset["x_test"]), 10000)
+        self.assertEqual(len(dataset["y_test"]), 10000)
 
     def test_cifar100_list(self):
         dataset = self.dataset_instance.get_list()
-        assert len(dataset) == 4
-        assert isinstance(dataset, list)
+        self.assertEqual(len(dataset), 4)
+        self.assertIsInstance(dataset, list)
diff --git a/tmu/lib/cpp/tests/MNISTDemoCPP.py b/tmu/lib/cpp/tests/MNISTDemoCPP.py
new file mode 100644
index 00000000..349a466b
--- /dev/null
+++ b/tmu/lib/cpp/tests/MNISTDemoCPP.py
@@ -0,0 +1,96 @@
+import sys
+import pathlib
+import subprocess
+import numpy as np
+
+try:
+    import tmulibpy
+except ImportError as e:
+    print("Failed to import tmulibpy. Attempting to import from manually built extension...")
+    current_dir = pathlib.Path(__file__).parent
+    cmake_dir = current_dir / ".." / ".." / "cmake-build-release"
+    root_dir = current_dir / ".." / ".." / ".." / ".."
+    sys.path.append(str(cmake_dir.absolute()))
+    import tmulibpy
+
+# Attempt to import tmu
+try:
+    import tmu
+    from tmu.models.classification.vanilla_classifier import TMClassifier
+except ImportError as e:
+    print(e)
+    print("Attempting to install the package in editable mode from the project's root directory...")
+    # Execute pip install command
+    try:
+        subprocess.check_call([sys.executable, "-m", "pip", "install", "-e", str(root_dir.absolute()), "-vvv"])
+    except subprocess.CalledProcessError as e:
+        print("Failed to install the package. Please run 'pip install -e .' manually.")
+
+
+import logging
+import argparse
+from tmu.data import MNIST
+from tmu.models.classification.vanilla_classifier import TMClassifier
+from tmu.tools import BenchmarkTimer
+from tmu.util.cuda_profiler import CudaProfiler
+
+_LOGGER = logging.getLogger(__name__)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--num_clauses", default=2000, type=int)
+    parser.add_argument("--T", default=5000, type=int)
+    parser.add_argument("--s", default=10.0, type=float)
+    parser.add_argument("--max_included_literals", default=32, type=int)
+    parser.add_argument("--device", default="CPU", type=str, choices=["CPU", "CPU_sparse", "CUDA"])
+    parser.add_argument("--weighted_clauses", default=True, type=bool)
+    parser.add_argument("--epochs", default=60, type=int)
+    args = parser.parse_args()
+
+    data = MNIST().get()
+
+    tm = tmulibpy.TMVanillaClassifier(
+        T=args.T,
+        s=args.s,
+        d=200.0,
+        number_of_clauses=args.num_clauses,
+        confidence_driven_updating=False,
+        weighted_clauses=True,
+        type_i_feedback=True,
+        type_ii_feedback=True,
+        type_iii_feedback=False,
+        type_i_ii_ratio=1.0,
+        max_included_literals=0,
+        boost_true_positive_feedback=True,
+        reuse_random_feedback=True,
+        patch_dim=None,
+        number_of_state_bits=8,
+        number_of_state_bits_ind=8,
+        batch_size=100,
+        incremental=True,
+        seed=42,
+    )
+
+
+    _LOGGER.info(f"Running {TMClassifier} for {args.epochs}")
+    for epoch in range(args.epochs):
+        benchmark_total = BenchmarkTimer(logger=None, text="Epoch Time")
+        with benchmark_total:
+            benchmark1 = BenchmarkTimer(logger=None, text="Training Time")
+            with benchmark1:
+                res = tm.fit(
+                    data["x_train"].astype(np.uint32),
+                    data["y_train"].astype(np.uint32),
+                    metrics=["update_p"],
+                )
+
+            # print(res)
+            benchmark2 = BenchmarkTimer(logger=None, text="Testing Time")
+            with benchmark2:
+                result = 100 * (tm.predict(data["x_test"]) == data["y_test"]).mean()
+
+            _LOGGER.info(f"Epoch: {epoch + 1}, Accuracy: {result:.2f}, Training Time: {benchmark1.elapsed():.2f}s, "
+                         f"Testing Time: {benchmark2.elapsed():.2f}s")
+
+        if args.device == "CUDA":
+            CudaProfiler().print_timings(benchmark=benchmark_total)
diff --git a/tmu/lib/cpp/tests/mnist_generator.py b/tmu/lib/cpp/tests/mnist_generator.py
new file mode 100644
index 00000000..2c4a8f08
--- /dev/null
+++ b/tmu/lib/cpp/tests/mnist_generator.py
@@ -0,0 +1,16 @@
+from tmu.data import MNIST
+import numpy as np
+data = MNIST().get()
+print(data.keys())
+
+x_train = data["x_train"]
+y_train = data["y_train"]
+x_test = data["x_test"]
+y_test = data["y_test"]
+
+print(x_train.shape, y_train.shape, x_test.shape, y_test.shape)
+# 2d, save as files whhere value is separated by space
+np.savetxt("x_train.txt", x_train, fmt="%d")
+np.savetxt("y_train.txt", y_train, fmt="%d")
+np.savetxt("x_test.txt", x_test, fmt="%d")
+np.savetxt("y_test.txt", y_test, fmt="%d")