Repository: cair/tmu (language: python)

Issue:
Added github action for building on different platforms and python ve…
…rsions

Added initial structure of a generic datasets interface.


Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index b6e47617..8aee671c 100644
--- a/.gitignore
+++ b/.gitignore
@@ -14,7 +14,6 @@ dist/
 downloads/
 eggs/
 .eggs/
-lib/
 lib64/
 parts/
 sdist/
@@ -127,3 +126,7 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+
+.idea/
+
+*.o
diff --git a/examples/autoencoder/__init__.py b/examples/autoencoder/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/examples/classification/README.md b/examples/classification/README.md
new file mode 100644
index 00000000..21bea55a
--- /dev/null
+++ b/examples/classification/README.md
@@ -0,0 +1,1 @@
+# Examples - Classification
\ No newline at end of file
diff --git a/examples/classification/__init__.py b/examples/classification/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/examples/experimental/__init__.py b/examples/experimental/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/examples/IMDbAutoEncoderDemo.py b/examples/experimental/autoencoder/IMDbAutoEncoderDemo.py
similarity index 98%
rename from examples/IMDbAutoEncoderDemo.py
rename to examples/experimental/autoencoder/IMDbAutoEncoderDemo.py
index 43b04ebd..3ef5a430 100644
--- a/examples/IMDbAutoEncoderDemo.py
+++ b/examples/experimental/autoencoder/IMDbAutoEncoderDemo.py
@@ -8,7 +8,7 @@
 from sklearn.metrics import f1_score
 from sklearn.metrics.pairwise import cosine_similarity
 
-from tmu.tsetlin_machine import TMAutoEncoder
+from tmu.models.autoencoder.autoencoder import TMAutoEncoder
 
 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']
 
@@ -29,7 +29,7 @@
 print("Number of clauses:", clauses)
 
 NUM_WORDS=10000
-INDEX_FROM=2 
+INDEX_FROM=2
 
 print("Downloading dataset...")
 
@@ -53,7 +53,7 @@
 	terms = []
 	for word_id in train_x[i]:
 		terms.append(id_to_word[word_id].lower())
-	
+
 	training_documents.append(terms)
 
 testing_documents = []
@@ -61,7 +61,7 @@
 	terms = []
 	for word_id in test_x[i]:
 		terms.append(id_to_word[word_id].lower())
-	
+
 	testing_documents.append(terms)
 
 def tokenizer(s):
@@ -88,14 +88,14 @@ def tokenizer(s):
 	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(len(target_words)):
 		precision.append(tm.clause_precision(i, True, X_train, number_of_examples=500))
-	
+
 	print("Calculating recall\n")
 	recall = []
 	for i in range(len(target_words)):
@@ -123,7 +123,7 @@ def tokenizer(s):
 		profile[i,:] = np.where(weights >= clause_weight_threshold, weights, 0)
 
 	similarity = cosine_similarity(profile)
-	
+
 	print("\nWord Similarity\n")
 
 	for i in range(len(target_words)):
diff --git a/examples/experimental/autoencoder/__init__.py b/examples/experimental/autoencoder/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/examples/BillionWordsPredictionDemo.py b/examples/experimental/classification/BillionWordsPredictionDemo.py
similarity index 97%
rename from examples/BillionWordsPredictionDemo.py
rename to examples/experimental/classification/BillionWordsPredictionDemo.py
index 2fa5e00f..4084391f 100644
--- a/examples/BillionWordsPredictionDemo.py
+++ b/examples/experimental/classification/BillionWordsPredictionDemo.py
@@ -6,7 +6,7 @@
 from time import time
 from sklearn.feature_extraction.text import CountVectorizer
 from sklearn.model_selection import train_test_split
-from tmu.tsetlin_machine import TMClassifier
+from tmu.models.classification.vanilla_classifier import TMClassifier
 import pickle
 
 #target_word = 'awful'#'comedy'#'romance'#"scary"
@@ -30,7 +30,7 @@
 s = 5.0
 
 NUM_WORDS=10000
-INDEX_FROM=2 
+INDEX_FROM=2
 
 # Data obtained from https://www.kaggle.com/c/billion-word-imputation
 
@@ -64,7 +64,7 @@
 target_id = vectorizer_X.vocabulary_[target_word]
 Y = X_csc[:,target_id].toarray().reshape(X_csc.shape[0])
 cols = np.arange(X_csc.shape[1]) != target_id
-X_csc = X_csc[:,cols] 
+X_csc = X_csc[:,cols]
 
 X_csr = X_csc.tocsr()
 
@@ -140,7 +140,7 @@
 	for k in sorted_literals:
 		if literal_importance[k] == 0:
 			break
-		
+
 		print("¬" + feature_names[k - number_of_features], end=' ')
 
 	print()
@@ -160,4 +160,4 @@
 			break
 
 		print("¬" + feature_names[k - number_of_features], end=' ')
-	print()
\ No newline at end of file
+	print()
diff --git a/examples/CIFAR100CoalescedDemo3x3.py b/examples/experimental/classification/CIFAR100CoalescedDemo3x3.py
similarity index 96%
rename from examples/CIFAR100CoalescedDemo3x3.py
rename to examples/experimental/classification/CIFAR100CoalescedDemo3x3.py
index d12fc9b5..2e9e43f7 100644
--- a/examples/CIFAR100CoalescedDemo3x3.py
+++ b/examples/experimental/classification/CIFAR100CoalescedDemo3x3.py
@@ -1,4 +1,4 @@
-from tmu.tsetlin_machine import TMCoalescedClassifier
+from tmu.models.classification.coalesced_classifier import TMCoalescedClassifier
 import numpy as np
 from time import time
 from sklearn.metrics import f1_score
@@ -10,7 +10,7 @@
 
 clauses = 8000
 T = int(clauses*0.75)
-s = 10.0 
+s = 10.0
 patch_size = 3
 resolution = 8
 number_of_state_bits_ta = 10
@@ -26,11 +26,11 @@
 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) 
+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) 
+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)
 
diff --git a/examples/CIFAR100CoalescedDemo3x3LiteralBudget.py b/examples/experimental/classification/CIFAR100CoalescedDemo3x3LiteralBudget.py
similarity index 96%
rename from examples/CIFAR100CoalescedDemo3x3LiteralBudget.py
rename to examples/experimental/classification/CIFAR100CoalescedDemo3x3LiteralBudget.py
index 2d06aa9d..4b647dfe 100644
--- a/examples/CIFAR100CoalescedDemo3x3LiteralBudget.py
+++ b/examples/experimental/classification/CIFAR100CoalescedDemo3x3LiteralBudget.py
@@ -1,4 +1,4 @@
-from tmu.tsetlin_machine import TMCoalescedClassifier
+from tmu.models.classification.coalesced_classifier import TMCoalescedClassifier
 import numpy as np
 from time import time
 from sklearn.metrics import f1_score
@@ -11,7 +11,7 @@
 max_included_literals=32
 clauses = 80000
 T = int(clauses*0.75)
-s = 10.0 
+s = 10.0
 patch_size = 3
 resolution = 8
 number_of_state_bits_ta = 8
@@ -27,11 +27,11 @@
 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) 
+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) 
+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)
 
diff --git a/examples/CIFAR100Demo3x3.py b/examples/experimental/classification/CIFAR100Demo3x3.py
similarity index 95%
rename from examples/CIFAR100Demo3x3.py
rename to examples/experimental/classification/CIFAR100Demo3x3.py
index 9e61332d..a8cad66f 100644
--- a/examples/CIFAR100Demo3x3.py
+++ b/examples/experimental/classification/CIFAR100Demo3x3.py
@@ -1,4 +1,4 @@
-from tmu.tsetlin_machine import TMCoalescedClassifier
+from tmu.models.classification.coalesced_classifier import TMCoalescedClassifier
 import numpy as np
 from time import time
 
@@ -9,7 +9,7 @@
 
 clauses = 80000
 T = int(clauses//10*0.75)
-s = 10.0 
+s = 10.0
 patch_size = 3
 resolution = 8
 number_of_state_bits_ta = 10
@@ -23,11 +23,11 @@
 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) 
+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) 
+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)
 
diff --git a/examples/CIFAR100MultiTaskDemo3x3.py b/examples/experimental/classification/CIFAR100MultiTaskDemo3x3.py
similarity index 96%
rename from examples/CIFAR100MultiTaskDemo3x3.py
rename to examples/experimental/classification/CIFAR100MultiTaskDemo3x3.py
index c581c930..90848139 100644
--- a/examples/CIFAR100MultiTaskDemo3x3.py
+++ b/examples/experimental/classification/CIFAR100MultiTaskDemo3x3.py
@@ -1,4 +1,4 @@
-from tmu.tsetlin_machine import TMMultiTaskClassifier
+from tmu.models.classification.multitask_classifier import TMMultiTaskClassifier
 import numpy as np
 from time import time
 from sklearn.metrics import f1_score
@@ -10,7 +10,7 @@
 
 clauses = 8000
 T = int(clauses*0.75)
-s = 10.0 
+s = 10.0
 patch_size = 3
 resolution = 8
 number_of_state_bits_ta = 10
@@ -28,11 +28,11 @@
 
 print((Y_train==1).sum())
 
-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_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) 
+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)
 
diff --git a/examples/CIFAR100MultiTaskSampleDemo3x3.py b/examples/experimental/classification/CIFAR100MultiTaskSampleDemo3x3.py
similarity index 97%
rename from examples/CIFAR100MultiTaskSampleDemo3x3.py
rename to examples/experimental/classification/CIFAR100MultiTaskSampleDemo3x3.py
index 6d9604f6..b59e6da1 100644
--- a/examples/CIFAR100MultiTaskSampleDemo3x3.py
+++ b/examples/experimental/classification/CIFAR100MultiTaskSampleDemo3x3.py
@@ -1,4 +1,4 @@
-from tmu.tsetlin_machine import TMMultiTaskClassifier
+from tmu.models.classification.multitask_classifier import TMMultiTaskClassifier
 import numpy as np
 from time import time
 from sklearn.metrics import f1_score
@@ -10,7 +10,7 @@
 
 clauses = 8000
 T = int(clauses*0.75)
-s = 10.0 
+s = 10.0
 patch_size = 3
 resolution = 8
 number_of_state_bits_ta = 10
@@ -28,11 +28,11 @@
 
 print((Y_train==1).sum())
 
-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_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) 
+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)
 
diff --git a/examples/CIFAR10CoalescedDemo3x3.py b/examples/experimental/classification/CIFAR10CoalescedDemo3x3.py
similarity index 96%
rename from examples/CIFAR10CoalescedDemo3x3.py
rename to examples/experimental/classification/CIFAR10CoalescedDemo3x3.py
index eb118676..3463eb4e 100644
--- a/examples/CIFAR10CoalescedDemo3x3.py
+++ b/examples/experimental/classification/CIFAR10CoalescedDemo3x3.py
@@ -1,4 +1,4 @@
-from tmu.tsetlin_machine import TMCoalescedClassifier
+from tmu.models.classification.coalesced_classifier import TMCoalescedClassifier
 import numpy as np
 from time import time
 from sklearn.metrics import f1_score
@@ -10,7 +10,7 @@
 
 clauses = 8000
 T = int(clauses*0.75)
-s = 10.0 
+s = 10.0
 patch_size = 3
 resolution = 8
 number_of_state_bits_ta = 8
@@ -26,11 +26,11 @@
 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) 
+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) 
+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)
 
diff --git a/examples/CIFAR10CoalescedDemo3x3LiteralBudget.py b/examples/experimental/classification/CIFAR10CoalescedDemo3x3LiteralBudget.py
similarity index 96%
rename from examples/CIFAR10CoalescedDemo3x3LiteralBudget.py
rename to examples/experimental/classification/CIFAR10CoalescedDemo3x3LiteralBudget.py
index 12127262..9553e259 100644
--- a/examples/CIFAR10CoalescedDemo3x3LiteralBudget.py
+++ b/examples/experimental/classification/CIFAR10CoalescedDemo3x3LiteralBudget.py
@@ -1,4 +1,4 @@
-from tmu.tsetlin_machine import TMCoalescedClassifier
+from tmu.models.classification.coalesced_classifier import TMCoalescedClassifier
 import numpy as np
 from time import time
 from sklearn.metrics import f1_score
@@ -11,7 +11,7 @@
 max_included_literals=32
 clauses = 80000
 T = int(clauses*0.75)
-s = 10.0 
+s = 10.0
 patch_size = 3
 resolution = 8
 number_of_state_bits_ta = 8
@@ -27,11 +27,11 @@
 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) 
+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) 
+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)
 
diff --git a/examples/CIFAR10MultiTaskSampleDemo3x3.py b/examples/experimental/classification/CIFAR10MultiTaskSampleDemo3x3.py
similarity index 97%
rename from examples/CIFAR10MultiTaskSampleDemo3x3.py
rename to examples/experimental/classification/CIFAR10MultiTaskSampleDemo3x3.py
index 4263d121..8522cee0 100644
--- a/examples/CIFAR10MultiTaskSampleDemo3x3.py
+++ b/examples/experimental/classification/CIFAR10MultiTaskSampleDemo3x3.py
@@ -1,4 +1,4 @@
-from tmu.tsetlin_machine import TMMultiTaskClassifier
+from tmu.models.classification.multitask_classifier import TMMultiTaskClassifier
 import numpy as np
 from time import time
 from sklearn.metrics import f1_score
@@ -10,7 +10,7 @@
 
 clauses = 80000
 T = int(clauses*0.75)
-s = 10.0 
+s = 10.0
 patch_size = 3
 resolution = 8
 number_of_state_bits_ta = 8
@@ -28,11 +28,11 @@
 
 print((Y_train==1).sum())
 
-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_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) 
+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)
 
diff --git a/examples/CIFAR2Demo3x3LiteralBudget.py b/examples/experimental/classification/CIFAR2Demo3x3LiteralBudget.py
similarity index 96%
rename from examples/CIFAR2Demo3x3LiteralBudget.py
rename to examples/experimental/classification/CIFAR2Demo3x3LiteralBudget.py
index e2a593c3..5f6c7945 100644
--- a/examples/CIFAR2Demo3x3LiteralBudget.py
+++ b/examples/experimental/classification/CIFAR2Demo3x3LiteralBudget.py
@@ -1,4 +1,4 @@
-from tmu.tsetlin_machine import TMClassifier
+from tmu.models.classification.vanilla_classifier import TMClassifier
 import numpy as np
 from time import time
 from sklearn.metrics import f1_score
@@ -14,7 +14,7 @@
 max_included_literals=32
 clauses = 8000
 T = int(clauses*0.75)
-s = 10.0 
+s = 10.0
 patch_size = 3
 resolution = 8
 number_of_state_bits_ta = 8
@@ -30,11 +30,11 @@
 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) 
+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) 
+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)
 
diff --git a/examples/CIFARDemo3x3.py b/examples/experimental/classification/CIFARDemo3x3.py
similarity index 95%
rename from examples/CIFARDemo3x3.py
rename to examples/experimental/classification/CIFARDemo3x3.py
index 286620ff..70403842 100644
--- a/examples/CIFARDemo3x3.py
+++ b/examples/experimental/classification/CIFARDemo3x3.py
@@ -1,4 +1,4 @@
-from tmu.tsetlin_machine import TMCoalescedClassifier
+from tmu.models.classification.coalesced_classifier import TMCoalescedClassifier
 import numpy as np
 from time import time
 
@@ -8,7 +8,7 @@
 
 clauses = 2500
 T = int(clauses*0.75)
-s = 5.0 
+s = 5.0
 patch_size = 3
 resolution = 8
 number_of_state_bits = 8
@@ -25,11 +25,11 @@
 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) 
+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) 
+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)
 
@@ -54,4 +54,4 @@
                 print("%d %d %.2f %.2f %.2f %.2f" % (e, i, result_train, result_test, stop_training-start_training, stop_testing-start_testing))
                 print("%d %d %.2f %.2f %.2f %.2f" % (e, i, result_train, result_test, stop_training-start_training, stop_testing-start_testing), file=f)
                 f.flush()
-f.close()
\ No newline at end of file
+f.close()
diff --git a/examples/CIFARDemoVanilla3x3.py b/examples/experimental/classification/CIFARDemoVanilla3x3.py
similarity index 95%
rename from examples/CIFARDemoVanilla3x3.py
rename to examples/experimental/classification/CIFARDemoVanilla3x3.py
index 0561ff9e..1d56ebbe 100644
--- a/examples/CIFARDemoVanilla3x3.py
+++ b/examples/experimental/classification/CIFARDemoVanilla3x3.py
@@ -1,4 +1,4 @@
-from tmu.tsetlin_machine import TMClassifier
+from tmu.models.classification.vanilla_classifier import TMClassifier
 import numpy as np
 from time import time
 
@@ -8,7 +8,7 @@
 
 clauses = 128
 T = int(clauses*0.75)
-s = 10.0 
+s = 10.0
 patch_size = 3
 resolution = 8
 number_of_state_bits = 8
@@ -25,11 +25,11 @@
 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) 
+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) 
+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)
 
@@ -54,4 +54,4 @@
                 print("%d %d %.2f %.2f %d %.2f %.2f" % (e, i, result_train, result_test, tm.literal_clause_frequency().sum(), stop_training-start_training, stop_testing-start_testing))
                 print("%d %d %.2f %.2f %d %.2f %.2f" % (e, i, result_train, result_test, tm.literal_clause_frequency().sum(), stop_training-start_training, stop_testing-start_testing), file=f)
                 f.flush()
-f.close()
\ No newline at end of file
+f.close()
diff --git a/examples/FashionMNISTDemo3x3.py b/examples/experimental/classification/FashionMNISTDemo3x3.py
similarity index 90%
rename from examples/FashionMNISTDemo3x3.py
rename to examples/experimental/classification/FashionMNISTDemo3x3.py
index 2d4bc162..bdb7e6e3 100644
--- a/examples/FashionMNISTDemo3x3.py
+++ b/examples/experimental/classification/FashionMNISTDemo3x3.py
@@ -1,4 +1,4 @@
-from tmu.tsetlin_machine import TMCoalescedClassifier
+from tmu.models.classification.coalesced_classifier import TMCoalescedClassifier
 import numpy as np
 from time import time
 
@@ -21,11 +21,11 @@
 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], resolution), dtype=np.uint8) 
+X_train = np.empty((X_train_org.shape[0], X_train_org.shape[1], X_train_org.shape[2], 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], resolution), dtype=np.uint8) 
+X_test = np.empty((X_test_org.shape[0], X_test_org.shape[1], X_test_org.shape[2], resolution), dtype=np.uint8)
 for z in range(resolution):
         X_test[:,:,:,z] = X_test_org[:,:,:] >= (z+1)*255/(resolution+1)
 
diff --git a/examples/IMDbMultiWordPredictionDemo.py b/examples/experimental/classification/IMDbMultiWordPredictionDemo.py
similarity index 98%
rename from examples/IMDbMultiWordPredictionDemo.py
rename to examples/experimental/classification/IMDbMultiWordPredictionDemo.py
index 05d2102f..4b833b85 100644
--- a/examples/IMDbMultiWordPredictionDemo.py
+++ b/examples/experimental/classification/IMDbMultiWordPredictionDemo.py
@@ -6,7 +6,7 @@
 from time import time
 from sklearn.feature_extraction.text import CountVectorizer
 
-from tmu.tsetlin_machine import TMClassifier
+from tmu.models.classification.vanilla_classifier import TMClassifier
 
 #target_words = ['masterpiece', 'brilliant', 'comedy', 'scary', 'funny', 'hate', 'love', 'awful', 'terrible']
 
@@ -27,7 +27,7 @@
 s = 5.0
 
 NUM_WORDS=10000
-INDEX_FROM=2 
+INDEX_FROM=2
 
 print("Downloading dataset...")
 
@@ -51,7 +51,7 @@
 	terms = []
 	for word_id in train_x[i]:
 		terms.append(id_to_word[word_id].lower())
-	
+
 	training_documents.append(terms)
 
 testing_documents = []
@@ -59,7 +59,7 @@
 	terms = []
 	for word_id in test_x[i]:
 		terms.append(id_to_word[word_id].lower())
-	
+
 	testing_documents.append(terms)
 
 def tokenizer(s):
@@ -131,7 +131,7 @@ def tokenizer(s):
 		sorted_literals = np.argsort(-1*literal_importance)[0:profile_size]
 		for k in sorted_literals:
 			if literal_importance[k] == 0:
-				break			
+				break
 			print("¬" + feature_names[k - number_of_features], end=' ')
 
 		print()
diff --git a/examples/IMDbMultiWordPredictionDemoCoalesced.py b/examples/experimental/classification/IMDbMultiWordPredictionDemoCoalesced.py
similarity index 97%
rename from examples/IMDbMultiWordPredictionDemoCoalesced.py
rename to examples/experimental/classification/IMDbMultiWordPredictionDemoCoalesced.py
index 755ea71e..f35ac00b 100644
--- a/examples/IMDbMultiWordPredictionDemoCoalesced.py
+++ b/examples/experimental/classification/IMDbMultiWordPredictionDemoCoalesced.py
@@ -6,7 +6,7 @@
 from time import time
 from sklearn.feature_extraction.text import CountVectorizer
 
-from tmu.tsetlin_machine import TMCoalescedClassifier
+from tmu.models.classification.coalesced_classifier import TMCoalescedClassifier
 
 #target_words = ['masterpiece', 'brilliant', 'comedy', 'scary', 'funny', 'hate', 'love', 'awful', 'terrible']
 
@@ -25,7 +25,7 @@
 s = 5.0
 
 NUM_WORDS=10000
-INDEX_FROM=2 
+INDEX_FROM=2
 
 print("Downloading dataset...")
 
@@ -49,7 +49,7 @@
 	terms = []
 	for word_id in train_x[i]:
 		terms.append(id_to_word[word_id].lower())
-	
+
 	training_documents.append(terms)
 
 testing_documents = []
@@ -57,7 +57,7 @@
 	terms = []
 	for word_id in test_x[i]:
 		terms.append(id_to_word[word_id].lower())
-	
+
 	testing_documents.append(terms)
 
 def tokenizer(s):
diff --git a/examples/IMDbTextCategorizationDemo.py b/examples/experimental/classification/IMDbTextCategorizationDemo.py
similarity index 93%
rename from examples/IMDbTextCategorizationDemo.py
rename to examples/experimental/classification/IMDbTextCategorizationDemo.py
index e9b81da3..3e765c11 100644
--- a/examples/IMDbTextCategorizationDemo.py
+++ b/examples/experimental/classification/IMDbTextCategorizationDemo.py
@@ -5,12 +5,12 @@
 from keras.datasets import imdb
 from time import time
 
-from tmu.tsetlin_machine import TMClassifier
+from tmu.models.classification.vanilla_classifier import TMClassifier
 
 MAX_NGRAM = 2
 
 NUM_WORDS=5000
-INDEX_FROM=2 
+INDEX_FROM=2
 
 FEATURES=5000
 
@@ -38,18 +38,18 @@
 	terms = []
 	for word_id in train_x[i]:
 		terms.append(id_to_word[word_id])
-	
+
 	for N in range(1,MAX_NGRAM+1):
 		grams = [terms[j:j+N] for j in range(len(terms)-N+1)]
 		for gram in grams:
 			phrase = " ".join(gram)
-			
+
 			if phrase in vocabulary:
 				vocabulary[phrase] += 1
 			else:
 				vocabulary[phrase] = 1
 
-# Assign a bit position to each N-gram (minimum frequency 10) 
+# Assign a bit position to each N-gram (minimum frequency 10)
 
 phrase_bit_nr = {}
 bit_nr_phrase = {}
@@ -93,7 +93,7 @@
 		for gram in grams:
 			phrase = " ".join(gram)
 			if phrase in phrase_bit_nr:
-				X_test[i,phrase_bit_nr[phrase]] = 1				
+				X_test[i,phrase_bit_nr[phrase]] = 1
 
 	Y_test[i] = test_y[i]
 
diff --git a/examples/IMDbTextCategorizationDemoTypeIIIFeedback.py b/examples/experimental/classification/IMDbTextCategorizationDemoTypeIIIFeedback.py
similarity index 94%
rename from examples/IMDbTextCategorizationDemoTypeIIIFeedback.py
rename to examples/experimental/classification/IMDbTextCategorizationDemoTypeIIIFeedback.py
index 236ba092..a19ac83b 100644
--- a/examples/IMDbTextCategorizationDemoTypeIIIFeedback.py
+++ b/examples/experimental/classification/IMDbTextCategorizationDemoTypeIIIFeedback.py
@@ -5,12 +5,12 @@
 from keras.datasets import imdb
 from time import time
 
-from tmu.tsetlin_machine import TMClassifier
+from tmu.models.classification.vanilla_classifier import TMClassifier
 
 MAX_NGRAM = 2
 
 NUM_WORDS=5000
-INDEX_FROM=2 
+INDEX_FROM=2
 
 FEATURES=5000
 
@@ -38,18 +38,18 @@
 	terms = []
 	for word_id in train_x[i]:
 		terms.append(id_to_word[word_id])
-	
+
 	for N in range(1,MAX_NGRAM+1):
 		grams = [terms[j:j+N] for j in range(len(terms)-N+1)]
 		for gram in grams:
 			phrase = " ".join(gram)
-			
+
 			if phrase in vocabulary:
 				vocabulary[phrase] += 1
 			else:
 				vocabulary[phrase] = 1
 
-# Assign a bit position to each N-gram (minimum frequency 10) 
+# Assign a bit position to each N-gram (minimum frequency 10)
 
 phrase_bit_nr = {}
 bit_nr_phrase = {}
@@ -93,7 +93,7 @@
 		for gram in grams:
 			phrase = " ".join(gram)
 			if phrase in phrase_bit_nr:
-				X_test[i,phrase_bit_nr[phrase]] = 1				
+				X_test[i,phrase_bit_nr[phrase]] = 1
 
 	Y_test[i] = test_y[i]
 
diff --git a/examples/IMDbWordEmbeddingDemo.py b/examples/experimental/classification/IMDbWordEmbeddingDemo.py
similarity index 98%
rename from examples/IMDbWordEmbeddingDemo.py
rename to examples/experimental/classification/IMDbWordEmbeddingDemo.py
index 7dd3fdf5..100fb6df 100644
--- a/examples/IMDbWordEmbeddingDemo.py
+++ b/examples/experimental/classification/IMDbWordEmbeddingDemo.py
@@ -8,7 +8,7 @@
 from sklearn.metrics import f1_score
 from sklearn.metrics.pairwise import cosine_similarity
 
-from tmu.tsetlin_machine import TMAutoEncoder
+from tmu.models.autoencoder.autoencoder import TMAutoEncoder
 
 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']
 
@@ -25,7 +25,7 @@
 print("Number of clauses:", clauses)
 
 NUM_WORDS=10000
-INDEX_FROM=2 
+INDEX_FROM=2
 
 print("Downloading dataset...")
 
@@ -49,7 +49,7 @@
 	terms = []
 	for word_id in train_x[i]:
 		terms.append(id_to_word[word_id].lower())
-	
+
 	training_documents.append(terms)
 
 testing_documents = []
@@ -57,7 +57,7 @@
 	terms = []
 	for word_id in test_x[i]:
 		terms.append(id_to_word[word_id].lower())
-	
+
 	testing_documents.append(terms)
 
 def tokenizer(s):
@@ -84,14 +84,14 @@ def tokenizer(s):
 	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(len(target_words)):
 		precision.append(tm.clause_precision(i, True, X_train, number_of_examples=500))
-	
+
 	print("Calculating recall\n")
 	recall = []
 	for i in range(len(target_words)):
@@ -119,7 +119,7 @@ def tokenizer(s):
 		profile[i,:] = np.where(weights >= clause_weight_threshold, weights, 0)
 
 	similarity = cosine_similarity(profile)
-	
+
 	print("\nWord Similarity\n")
 
 	for i in range(len(target_words)):
diff --git a/examples/IMDbWordPredictionDemo.py b/examples/experimental/classification/IMDbWordPredictionDemo.py
similarity index 97%
rename from examples/IMDbWordPredictionDemo.py
rename to examples/experimental/classification/IMDbWordPredictionDemo.py
index 2be59ce3..cb65e123 100644
--- a/examples/IMDbWordPredictionDemo.py
+++ b/examples/experimental/classification/IMDbWordPredictionDemo.py
@@ -6,7 +6,7 @@
 from time import time
 from sklearn.feature_extraction.text import CountVectorizer
 
-from tmu.tsetlin_machine import TMClassifier
+from tmu.models.classification.vanilla_classifier import TMClassifier
 
 target_word = 'awful' #'frightening'#'comedy'#'romance'#"scary"
 #target_word = 'brilliant'
@@ -24,7 +24,7 @@
 print("Number of clauses:", clauses)
 
 NUM_WORDS=10000
-INDEX_FROM=2 
+INDEX_FROM=2
 
 print("Downloading dataset...")
 
@@ -48,7 +48,7 @@
 	terms = []
 	for word_id in train_x[i]:
 		terms.append(id_to_word[word_id].lower())
-	
+
 	training_documents.append(terms)
 
 testing_documents = []
@@ -56,7 +56,7 @@
 	terms = []
 	for word_id in test_x[i]:
 		terms.append(id_to_word[word_id].lower())
-	
+
 	testing_documents.append(terms)
 
 def tokenizer(s):
@@ -135,7 +135,7 @@ def tokenizer(s):
 	sorted_literals = np.argsort(-1*literal_importance)[0:profile_size]
 	for k in sorted_literals:
 		if literal_importance[k] == 0:
-			break		
+			break
 		print("¬" + feature_names[k - number_of_features], end=' ')
 
 	print()
@@ -155,4 +155,4 @@ def tokenizer(s):
 		print("¬" + feature_names[k - number_of_features], end=' ')
 
 	print()
-	print("\n#%d Accuracy: %.2f%% Training: %.2fs Testing: %.2fs" % (i+1, result, stop_training-start_training, stop_testing-start_testing))
\ No newline at end of file
+	print("\n#%d Accuracy: %.2f%% Training: %.2fs Testing: %.2fs" % (i+1, result, stop_training-start_training, stop_testing-start_testing))
diff --git a/examples/IMDbWordPredictionMultiTaskDemo.py b/examples/experimental/classification/IMDbWordPredictionMultiTaskDemo.py
similarity index 97%
rename from examples/IMDbWordPredictionMultiTaskDemo.py
rename to examples/experimental/classification/IMDbWordPredictionMultiTaskDemo.py
index 97ee36c1..78ef3c25 100644
--- a/examples/IMDbWordPredictionMultiTaskDemo.py
+++ b/examples/experimental/classification/IMDbWordPredictionMultiTaskDemo.py
@@ -8,7 +8,7 @@
 from sklearn.metrics.pairwise import cosine_similarity
 from scipy.sparse import csr_matrix, csc_matrix, lil_matrix
 
-from tmu.tsetlin_machine import TMMultiTaskClassifier
+from tmu.models.classification.multitask_classifier import TMMultiTaskClassifier
 
 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']
 
@@ -30,7 +30,7 @@
 print("Number of clauses:", clauses)
 
 NUM_WORDS=10000
-INDEX_FROM=2 
+INDEX_FROM=2
 
 print("Downloading dataset...")
 
@@ -54,7 +54,7 @@
 	terms = []
 	for word_id in train_x[i]:
 		terms.append(id_to_word[word_id].lower())
-	
+
 	training_documents.append(terms)
 
 testing_documents = []
@@ -62,7 +62,7 @@
 	terms = []
 	for word_id in test_x[i]:
 		terms.append(id_to_word[word_id].lower())
-	
+
 	testing_documents.append(terms)
 
 def tokenizer(s):
@@ -90,7 +90,7 @@ def tokenizer(s):
 			target_indices = np.setdiff1d(np.arange(X_train_csr.shape[0], dtype=np.uint32), X_train_csc[:,target_id].indices)
 		examples = np.random.choice(target_indices, size=context_size, replace=True)
 		X_train[i][e,:] = (X_train_csr[examples].toarray().sum(axis=0) > 0).astype(np.uint32)
-		X_train[i][e,target_id] = 1 
+		X_train[i][e,target_id] = 1
 		Y_train[i][e] = target
 	X_train[i] = X_train[i].tocsr()
 	Y_train[i] = Y_train[i].tocsr()
@@ -132,7 +132,7 @@ def tokenizer(s):
 	for i in range(len(target_words)):
 		result.append(100*(prediction[i] == Y_test[i]).mean())
 	stop_testing = time()
-	
+
 	for j in range(clauses):
 		print("Clause #%d " % (j), end=' ')
 		for i in range(len(target_words)):
@@ -153,7 +153,7 @@ def tokenizer(s):
 		profile[i,:] = np.where(weights >= clause_weight_threshold, weights, 0)
 
 	similarity = cosine_similarity(profile)
-	
+
 	print("\nWord Similarity\n")
 
 	for i in range(len(target_words)):
@@ -164,4 +164,4 @@ def tokenizer(s):
 		print()
 
 	print("\n#%d Accuracy: %s Training: %.2fs Testing: %.2fs" % (e+1, str(result), stop_training-start_training, stop_testing-start_testing))
-	print()
\ No newline at end of file
+	print()
diff --git a/examples/InterpretabilityDemo.py b/examples/experimental/classification/InterpretabilityDemo.py
similarity index 97%
rename from examples/InterpretabilityDemo.py
rename to examples/experimental/classification/InterpretabilityDemo.py
index 9354d42c..9d353fd8 100644
--- a/examples/InterpretabilityDemo.py
+++ b/examples/experimental/classification/InterpretabilityDemo.py
@@ -1,4 +1,4 @@
-from tmu.tsetlin_machine import TMClassifier
+from tmu.models.classification.vanilla_classifier import TMClassifier
 import numpy as np
 from time import time
 
diff --git a/examples/InterpretabilityDemoAND.py b/examples/experimental/classification/InterpretabilityDemoAND.py
similarity index 97%
rename from examples/InterpretabilityDemoAND.py
rename to examples/experimental/classification/InterpretabilityDemoAND.py
index ab808e85..8987d00d 100644
--- a/examples/InterpretabilityDemoAND.py
+++ b/examples/experimental/classification/InterpretabilityDemoAND.py
@@ -1,5 +1,5 @@
-from tmu.tsetlin_machine import TMMultiChannelClassifier
-import numpy as np 
+from tmu.models.classification.multichannel_classifier import TMMultiChannelClassifier
+import numpy as np
 
 number_of_and_components = 2
 clauses = 4
@@ -71,7 +71,7 @@
 
 		if X_test[0,e,0] == 1 and X_test[0,e,1] == 0:
 			count_10 += 1
-	
+
 		count += 1
 
 print("Y=0")
@@ -98,7 +98,7 @@
 
 		if X_test[0,e,0] == 1 and X_test[0,e,1] == 0:
 			count_10 += 1
-	
+
 		count += 1
 
 print("Y=1")
@@ -127,7 +127,7 @@
 
 		if X_train[0,e,0] == 1 and X_train[0,e,1] == 0:
 			count_10 += 1
-	
+
 		count += 1
 
 print("Y=0")
@@ -154,7 +154,7 @@
 
 		if X_train[0,e,0] == 1 and X_train[0,e,1] == 0:
 			count_10 += 1
-	
+
 		count += 1
 
 print("Y=1")
@@ -180,7 +180,7 @@
 
 for c in range(number_of_and_components):
 	print("Channel #%d" % (c))
-		
+
 	for j in range(clauses):
 		print("\tClause #%d " % (j), end=' ')
 		l = []
diff --git a/examples/InterpretabilityDemoCoalesced.py b/examples/experimental/classification/InterpretabilityDemoCoalesced.py
similarity index 95%
rename from examples/InterpretabilityDemoCoalesced.py
rename to examples/experimental/classification/InterpretabilityDemoCoalesced.py
index 07e546f5..3317603a 100644
--- a/examples/InterpretabilityDemoCoalesced.py
+++ b/examples/experimental/classification/InterpretabilityDemoCoalesced.py
@@ -1,4 +1,4 @@
-from tmu.tsetlin_machine import TMCoalescedClassifier
+from tmu.models.classification.coalesced_classifier import TMCoalescedClassifier
 import numpy as np
 from time import time
 
diff --git a/examples/InterpretabilityDemoOneVsOne.py b/examples/experimental/classification/InterpretabilityDemoOneVsOne.py
similarity index 95%
rename from examples/InterpretabilityDemoOneVsOne.py
rename to examples/experimental/classification/InterpretabilityDemoOneVsOne.py
index dbc755a3..16a0fe8f 100644
--- a/examples/InterpretabilityDemoOneVsOne.py
+++ b/examples/experimental/classification/InterpretabilityDemoOneVsOne.py
@@ -1,4 +1,4 @@
-from tmu.tsetlin_machine import TMOneVsOneClassifier
+from tmu.models.classification.one_vs_one_classifier import TMOneVsOneClassifier
 import numpy as np
 from time import time
 
diff --git a/examples/MNISTClauseCoOccurrenceDemo.py b/examples/experimental/classification/MNISTClauseCoOccurrenceDemo.py
similarity index 90%
rename from examples/MNISTClauseCoOccurrenceDemo.py
rename to examples/experimental/classification/MNISTClauseCoOccurrenceDemo.py
index faf69d74..71b38462 100644
--- a/examples/MNISTClauseCoOccurrenceDemo.py
+++ b/examples/experimental/classification/MNISTClauseCoOccurrenceDemo.py
@@ -3,7 +3,7 @@
 
 from keras.datasets import mnist
 
-from tmu.tsetlin_machine import TMClassifier
+from tmu.models.classification.vanilla_classifier import TMClassifier
 
 np.set_printoptions(threshold=np.inf, linewidth=250, precision=2, suppress=True)
 
@@ -19,11 +19,11 @@
 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], resolution), dtype=np.uint8) 
+X_train = np.empty((X_train_org.shape[0], X_train_org.shape[1], X_train_org.shape[2], 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], resolution), dtype=np.uint8) 
+X_test = np.empty((X_test_org.shape[0], X_test_org.shape[1], X_test_org.shape[2], resolution), dtype=np.uint8)
 for z in range(resolution):
         X_test[:,:,:,z] = X_test_org[:,:,:] >= (z+1)*255/(resolution+1)
 
@@ -49,4 +49,4 @@
 	print("#%d Accuracy: %.2f%% Training: %.2fs Testing: %.2fs" % (i+1, result, stop_training-start_training, stop_testing-start_testing))
 
 print()
-print(tm.clause_co_occurrence(X_test, percentage=True).toarray())
\ No newline at end of file
+print(tm.clause_co_occurrence(X_test, percentage=True).toarray())
diff --git a/examples/MNISTConvolutionLiteralBudgetDemo.py b/examples/experimental/classification/MNISTConvolutionLiteralBudgetDemo.py
similarity index 92%
rename from examples/MNISTConvolutionLiteralBudgetDemo.py
rename to examples/experimental/classification/MNISTConvolutionLiteralBudgetDemo.py
index a47e7692..6fd8dcc2 100644
--- a/examples/MNISTConvolutionLiteralBudgetDemo.py
+++ b/examples/experimental/classification/MNISTConvolutionLiteralBudgetDemo.py
@@ -3,7 +3,7 @@
 
 from keras.datasets import mnist
 
-from tmu.tsetlin_machine import TMClassifier
+from tmu.models.classification.vanilla_classifier import TMClassifier
 
 ensembles = 5
 epochs = 250
@@ -15,7 +15,7 @@
 
 (X_train, Y_train), (X_test, Y_test) = mnist.load_data()
 
-X_train = np.where(X_train >= 75, 1, 0) 
+X_train = np.where(X_train >= 75, 1, 0)
 X_test = np.where(X_test >= 75, 1, 0)
 
 f = open("mnist_convolution_%.1f_%d_%d_%d.txt" % (s, clauses, T, max_included_literals), "w+")
@@ -42,4 +42,4 @@
                 print("%d %d %.2f %.2f %.2f %.2f" % (ensemble, epoch, number_of_includes, result_test, stop_training-start_training, stop_testing-start_testing))
                 print("%d %d %.2f %.2f %.2f %.2f" % (ensemble, epoch, number_of_includes, result_test, stop_training-start_training, stop_testing-start_testing), file=f)
                 f.flush()
-f.close()
\ No newline at end of file
+f.close()
diff --git a/examples/MNISTDemo.py b/examples/experimental/classification/MNISTDemo.py
similarity index 88%
rename from examples/MNISTDemo.py
rename to examples/experimental/classification/MNISTDemo.py
index a93a2813..ec8ee457 100644
--- a/examples/MNISTDemo.py
+++ b/examples/experimental/classification/MNISTDemo.py
@@ -3,12 +3,12 @@
 
 from keras.datasets import mnist
 
-from tmu.tsetlin_machine import TMClassifier
+from tmu.models.classification.vanilla_classifier import TMClassifier
 
 (X_train, Y_train), (X_test, Y_test) = mnist.load_data()
 
-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) 
+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, 50, 10.0)
 
diff --git a/examples/MNISTDemo2DConvolutionOneVsOne.py b/examples/experimental/classification/MNISTDemo2DConvolutionOneVsOne.py
similarity index 73%
rename from examples/MNISTDemo2DConvolutionOneVsOne.py
rename to examples/experimental/classification/MNISTDemo2DConvolutionOneVsOne.py
index 3b636097..14251a47 100644
--- a/examples/MNISTDemo2DConvolutionOneVsOne.py
+++ b/examples/experimental/classification/MNISTDemo2DConvolutionOneVsOne.py
@@ -3,12 +3,12 @@
 
 from keras.datasets import mnist
 
-from tmu.tsetlin_machine import TMOneVsOneClassifier
+from tmu.models.classification.one_vs_one_classifier import TMOneVsOneClassifier
 
 (X_train, Y_train), (X_test, Y_test) = mnist.load_data()
 
-X_train = np.where(X_train >= 75, 1, 0) 
-X_test = np.where(X_test >= 75, 1, 0) 
+X_train = np.where(X_train >= 75, 1, 0)
+X_test = np.where(X_test >= 75, 1, 0)
 
 tm = TMOneVsOneClassifier(20000, 1500, 5.0, patch_dim=(10, 10), weighted_clauses=True)
 
@@ -17,7 +17,7 @@
 	start = time()
 	tm.fit(X_train, Y_train)
 	stop = time()
-	
+
 	result = 100*(tm.predict(X_test) == Y_test).mean()
-	
+
 	print("#%d Accuracy: %.2f%% (%.2fs)" % (i+1, result, stop-start))
diff --git a/examples/MNISTDemo2DConvolutionWeightedClauses.py b/examples/experimental/classification/MNISTDemo2DConvolutionWeightedClauses.py
similarity index 80%
rename from examples/MNISTDemo2DConvolutionWeightedClauses.py
rename to examples/experimental/classification/MNISTDemo2DConvolutionWeightedClauses.py
index 0a2e2901..1cd4ee9c 100644
--- a/examples/MNISTDemo2DConvolutionWeightedClauses.py
+++ b/examples/experimental/classification/MNISTDemo2DConvolutionWeightedClauses.py
@@ -3,12 +3,12 @@
 
 from keras.datasets import mnist
 
-from tmu.tsetlin_machine import TMClassifier
+from tmu.models.classification.vanilla_classifier import TMClassifier
 
 (X_train, Y_train), (X_test, Y_test) = mnist.load_data()
 
-X_train = np.where(X_train >= 75, 1, 0) 
-X_test = np.where(X_test >= 75, 1, 0) 
+X_train = np.where(X_train >= 75, 1, 0)
+X_test = np.where(X_test >= 75, 1, 0)
 
 tm = TMClassifier(2000, int(50*100), 5.0, platform='CPU', patch_dim=(10, 10), weighted_clauses=True)
 
@@ -17,7 +17,7 @@
 	start_training = time()
 	tm.fit(X_train, Y_train)
 	stop_training = time()
-	
+
 	start_testing = time()
 	result = 100*(tm.predict(X_test) == Y_test).mean()
 	stop_testing = time()
diff --git a/examples/MNISTDemo2DConvolutionWeightedClausesFocusedNegativeSampling.py b/examples/experimental/classification/MNISTDemo2DConvolutionWeightedClausesFocusedNegativeSampling.py
similarity index 74%
rename from examples/MNISTDemo2DConvolutionWeightedClausesFocusedNegativeSampling.py
rename to examples/experimental/classification/MNISTDemo2DConvolutionWeightedClausesFocusedNegativeSampling.py
index 665f3aab..b7e58fdb 100644
--- a/examples/MNISTDemo2DConvolutionWeightedClausesFocusedNegativeSampling.py
+++ b/examples/experimental/classification/MNISTDemo2DConvolutionWeightedClausesFocusedNegativeSampling.py
@@ -3,12 +3,12 @@
 
 from keras.datasets import mnist
 
-from tmu.tsetlin_machine import TMCoalescedClassifier
+from tmu.models.classification.coalesced_classifier import TMCoalescedClassifier
 
 (X_train, Y_train), (X_test, Y_test) = mnist.load_data()
 
-X_train = np.where(X_train >= 75, 1, 0) 
-X_test = np.where(X_test >= 75, 1, 0) 
+X_train = np.where(X_train >= 75, 1, 0)
+X_test = np.where(X_test >= 75, 1, 0)
 
 tm = TMCoalescedClassifier(20000, int(50*100), 5.0, platform='CUDA', patch_dim=(10, 10), focused_negative_sampling=True, weighted_clauses=True)
 
@@ -17,9 +17,9 @@
 	start_training = time()
 	tm.fit(X_train, Y_train)
 	stop_training = time()
-	
+
 	start_testing = time()
 	result = 100*(tm.predict(X_test) == Y_test).mean()
 	stop_testing = time()
 
-	print("#%d Accuracy: %.2f%% Literals: %d Training: %.2fs Testing: %.2fs" % (i+1, result, tm.literal_clause_frequency().sum(), stop_training-start_training, stop_testing-start_testing))
\ No newline at end of file
+	print("#%d Accuracy: %.2f%% Literals: %d Training: %.2fs Testing: %.2fs" % (i+1, result, tm.literal_clause_frequency().sum(), stop_training-start_training, stop_testing-start_testing))
diff --git a/examples/MNISTDemo2DConvolutionWeightedClausesIncremental.py b/examples/experimental/classification/MNISTDemo2DConvolutionWeightedClausesIncremental.py
similarity index 80%
rename from examples/MNISTDemo2DConvolutionWeightedClausesIncremental.py
rename to examples/experimental/classification/MNISTDemo2DConvolutionWeightedClausesIncremental.py
index f96a4da7..553ec981 100644
--- a/examples/MNISTDemo2DConvolutionWeightedClausesIncremental.py
+++ b/examples/experimental/classification/MNISTDemo2DConvolutionWeightedClausesIncremental.py
@@ -3,12 +3,12 @@
 
 from keras.datasets import mnist
 
-from tmu.tsetlin_machine import TMClassifier
+from tmu.models.classification.vanilla_classifier import TMClassifier
 
 (X_train, Y_train), (X_test, Y_test) = mnist.load_data()
 
-X_train = np.where(X_train >= 75, 1, 0) 
-X_test = np.where(X_test >= 75, 1, 0) 
+X_train = np.where(X_train >= 75, 1, 0)
+X_test = np.where(X_test >= 75, 1, 0)
 
 tm = TMClassifier(2000, int(50*100), 5.0, platform='CPU', patch_dim=(10, 10), weighted_clauses=True, incremental=True)
 
@@ -17,7 +17,7 @@
 	start_training = time()
 	tm.fit(X_train, Y_train)
 	stop_training = time()
-	
+
 	start_testing = time()
 	result = 100*(tm.predict(X_test) == Y_test).mean()
 	stop_testing = time()
diff --git a/examples/MNISTDemoCoalesced.py b/examples/experimental/classification/MNISTDemoCoalesced.py
similarity index 88%
rename from examples/MNISTDemoCoalesced.py
rename to examples/experimental/classification/MNISTDemoCoalesced.py
index e8886808..faafe09b 100644
--- a/examples/MNISTDemoCoalesced.py
+++ b/examples/experimental/classification/MNISTDemoCoalesced.py
@@ -3,12 +3,12 @@
 
 from keras.datasets import mnist
 
-from tmu.tsetlin_machine import TMCoalescedClassifier
+from tmu.models.classification.coalesced_classifier import TMCoalescedClassifier
 
 (X_train, Y_train), (X_test, Y_test) = mnist.load_data()
 
-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) 
+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 = TMCoalescedClassifier(20000, 50*100, 10.0, weighted_clauses=True)
 
diff --git a/examples/MNISTDemoCoalescedFocusedNegativeSampling.py b/examples/experimental/classification/MNISTDemoCoalescedFocusedNegativeSampling.py
similarity index 89%
rename from examples/MNISTDemoCoalescedFocusedNegativeSampling.py
rename to examples/experimental/classification/MNISTDemoCoalescedFocusedNegativeSampling.py
index b513b002..c42782d5 100644
--- a/examples/MNISTDemoCoalescedFocusedNegativeSampling.py
+++ b/examples/experimental/classification/MNISTDemoCoalescedFocusedNegativeSampling.py
@@ -3,12 +3,12 @@
 
 from keras.datasets import mnist
 
-from tmu.tsetlin_machine import TMCoalescedClassifier
+from tmu.models.classification.coalesced_classifier import TMCoalescedClassifier
 
 (X_train, Y_train), (X_test, Y_test) = mnist.load_data()
 
-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) 
+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 = TMCoalescedClassifier(20000, 50*100, 10.0, platform='CUDA', focused_negative_sampling=True, weighted_clauses=True)
 
@@ -22,4 +22,4 @@
         result = 100*(tm.predict(X_test) == Y_test).mean()
         stop_testing = time()
 
-        print("#%d Accuracy: %.2f%% Literals: %d Training: %.2fs Testing: %.2fs" % (i+1, result, tm.literal_clause_frequency().sum(), stop_training-start_training, stop_testing-start_testing))
\ No newline at end of file
+        print("#%d Accuracy: %.2f%% Literals: %d Training: %.2fs Testing: %.2fs" % (i+1, result, tm.literal_clause_frequency().sum(), stop_training-start_training, stop_testing-start_testing))
diff --git a/examples/MNISTDemoCoalescedTypeIIIFeedback.py b/examples/experimental/classification/MNISTDemoCoalescedTypeIIIFeedback.py
similarity index 93%
rename from examples/MNISTDemoCoalescedTypeIIIFeedback.py
rename to examples/experimental/classification/MNISTDemoCoalescedTypeIIIFeedback.py
index fc9a6527..4ed2bc31 100644
--- a/examples/MNISTDemoCoalescedTypeIIIFeedback.py
+++ b/examples/experimental/classification/MNISTDemoCoalescedTypeIIIFeedback.py
@@ -1,7 +1,7 @@
 import numpy as np
 from time import time
 from keras.datasets import mnist
-from tmu.tsetlin_machine import TMCoalescedClassifier
+from tmu.models.classification.coalesced_classifier import TMCoalescedClassifier
 
 clauses = 20000
 T = 50*100
@@ -10,8 +10,8 @@
 
 (X_train, Y_train), (X_test, Y_test) = mnist.load_data()
 
-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) 
+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 = TMCoalescedClassifier(clauses, T, s, number_of_state_bits_ta=number_of_state_bits_ta, weighted_clauses=True, type_iii_feedback=True, focused_negative_sampling=True)
 
@@ -34,4 +34,4 @@
 
         print("%d %.2f %.2f %d %.2f %.2f" % (i, result_train, result_test, tm.literal_clause_frequency().sum(), stop_training-start_training, stop_testing-start_testing), file=f)
         f.flush()
-f.close()
\ No newline at end of file
+f.close()
diff --git a/examples/MNISTDemoOneVsOne.py b/examples/experimental/classification/MNISTDemoOneVsOne.py
similarity index 88%
rename from examples/MNISTDemoOneVsOne.py
rename to examples/experimental/classification/MNISTDemoOneVsOne.py
index 26c9b048..dd968c2f 100644
--- a/examples/MNISTDemoOneVsOne.py
+++ b/examples/experimental/classification/MNISTDemoOneVsOne.py
@@ -3,12 +3,12 @@
 
 from keras.datasets import mnist
 
-from tmu.tsetlin_machine import TMOneVsOneClassifier
+from tmu.models.classification.one_vs_one_classifier import TMOneVsOneClassifier
 
 (X_train, Y_train), (X_test, Y_test) = mnist.load_data()
 
-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) 
+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 = TMOneVsOneClassifier(20000, 1500, 10.0, weighted_clauses=True)
 
diff --git a/examples/MNISTDemoTrainLayerOne.py b/examples/experimental/classification/MNISTDemoTrainLayerOne.py
similarity index 91%
rename from examples/MNISTDemoTrainLayerOne.py
rename to examples/experimental/classification/MNISTDemoTrainLayerOne.py
index d6c0d838..12867bc2 100644
--- a/examples/MNISTDemoTrainLayerOne.py
+++ b/examples/experimental/classification/MNISTDemoTrainLayerOne.py
@@ -3,7 +3,7 @@
 
 from keras.datasets import mnist
 
-from tmu.tsetlin_machine import TMCoalescedClassifier
+from tmu.models.classification.coalesced_classifier import TMCoalescedClassifier
 
 import copy
 
@@ -19,11 +19,11 @@
 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], resolution), dtype=np.uint8) 
+X_train = np.empty((X_train_org.shape[0], X_train_org.shape[1], X_train_org.shape[2], 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], resolution), dtype=np.uint8) 
+X_test = np.empty((X_test_org.shape[0], X_test_org.shape[1], X_test_org.shape[2], resolution), dtype=np.uint8)
 for z in range(resolution):
         X_test[:,:,:,z] = X_test_org[:,:,:] >= (z+1)*255/(resolution+1)
 
@@ -37,7 +37,7 @@
 	start_training = time()
 	tm.fit(X_train, Y_train)
 	stop_training = time()
-	
+
 	start_testing = time()
 	result = 100*(tm.predict(X_test) == Y_test).mean()
 	stop_testing = time()
diff --git a/examples/MNISTDemoTypeIIIFeedback.py b/examples/experimental/classification/MNISTDemoTypeIIIFeedback.py
similarity index 93%
rename from examples/MNISTDemoTypeIIIFeedback.py
rename to examples/experimental/classification/MNISTDemoTypeIIIFeedback.py
index 93112f39..b204c58e 100644
--- a/examples/MNISTDemoTypeIIIFeedback.py
+++ b/examples/experimental/classification/MNISTDemoTypeIIIFeedback.py
@@ -3,7 +3,7 @@
 
 from keras.datasets import mnist
 
-from tmu.tsetlin_machine import TMClassifier
+from tmu.models.classification.vanilla_classifier import TMClassifier
 
 clauses = 2000
 T = 5000
@@ -12,8 +12,8 @@
 
 (X_train, Y_train), (X_test, Y_test) = mnist.load_data()
 
-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) 
+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(clauses, T, s, type_iii_feedback=True, number_of_state_bits_ta=number_of_state_bits_ta, number_of_state_bits_ind=8, platform='CPU', weighted_clauses=True)
 
@@ -36,4 +36,4 @@
 
         print("%d %.2f %.2f %d %.2f %.2f" % (i, result_train, result_test, tm.literal_clause_frequency().sum(), stop_training-start_training, stop_testing-start_testing), file=f)
         f.flush()
-f.close()
\ No newline at end of file
+f.close()
diff --git a/examples/MNISTDemoWeightedClauses.py b/examples/experimental/classification/MNISTDemoWeightedClauses.py
similarity index 89%
rename from examples/MNISTDemoWeightedClauses.py
rename to examples/experimental/classification/MNISTDemoWeightedClauses.py
index 823f3f6b..dfc4efad 100644
--- a/examples/MNISTDemoWeightedClauses.py
+++ b/examples/experimental/classification/MNISTDemoWeightedClauses.py
@@ -3,12 +3,12 @@
 
 from keras.datasets import mnist
 
-from tmu.tsetlin_machine import TMClassifier
+from tmu.models.classification.vanilla_classifier import TMClassifier
 
 (X_train, Y_train), (X_test, Y_test) = mnist.load_data()
 
-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) 
+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, 50*100, 10.0, platform='CUDA', weighted_clauses=True)
 
diff --git a/examples/MNISTDemoWeightedClausesIncremental.py b/examples/experimental/classification/MNISTDemoWeightedClausesIncremental.py
similarity index 89%
rename from examples/MNISTDemoWeightedClausesIncremental.py
rename to examples/experimental/classification/MNISTDemoWeightedClausesIncremental.py
index 25b37603..5a699976 100644
--- a/examples/MNISTDemoWeightedClausesIncremental.py
+++ b/examples/experimental/classification/MNISTDemoWeightedClausesIncremental.py
@@ -3,12 +3,12 @@
 
 from keras.datasets import mnist
 
-from tmu.tsetlin_machine import TMClassifier
+from tmu.models.classification.vanilla_classifier import TMClassifier
 
 (X_train, Y_train), (X_test, Y_test) = mnist.load_data()
 
-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) 
+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, 50*100, 10.0, platform='CPU', weighted_clauses=True, incremental=True)
 
diff --git a/examples/MNISTLiteralBudgetDemo.py b/examples/experimental/classification/MNISTLiteralBudgetDemo.py
similarity index 94%
rename from examples/MNISTLiteralBudgetDemo.py
rename to examples/experimental/classification/MNISTLiteralBudgetDemo.py
index 5d9aa9c4..5662ac2b 100644
--- a/examples/MNISTLiteralBudgetDemo.py
+++ b/examples/experimental/classification/MNISTLiteralBudgetDemo.py
@@ -3,7 +3,7 @@
 
 from keras.datasets import mnist
 
-from tmu.tsetlin_machine import TMClassifier
+from tmu.models.classification.vanilla_classifier import TMClassifier
 
 ensembles = 5
 epochs = 250
@@ -15,8 +15,8 @@
 
 (X_train, Y_train), (X_test, Y_test) = mnist.load_data()
 
-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) 
+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)
 
 f = open("mnist_%.1f_%d_%d_%d.txt" % (s, clauses, T, max_included_literals), "w+")
 for ensemble in range(ensembles):
@@ -42,4 +42,4 @@
                 print("%d %d %.2f %.2f %.2f %.2f" % (ensemble, epoch, number_of_includes, result_test, stop_training-start_training, stop_testing-start_testing))
                 print("%d %d %.2f %.2f %.2f %.2f" % (ensemble, epoch, number_of_includes, result_test, stop_training-start_training, stop_testing-start_testing), file=f)
                 f.flush()
-f.close()
\ No newline at end of file
+f.close()
diff --git a/examples/MNISTLiteralBudgetDemoCoalesced.py b/examples/experimental/classification/MNISTLiteralBudgetDemoCoalesced.py
similarity index 90%
rename from examples/MNISTLiteralBudgetDemoCoalesced.py
rename to examples/experimental/classification/MNISTLiteralBudgetDemoCoalesced.py
index 6924ac0c..da5f4775 100644
--- a/examples/MNISTLiteralBudgetDemoCoalesced.py
+++ b/examples/experimental/classification/MNISTLiteralBudgetDemoCoalesced.py
@@ -3,12 +3,12 @@
 
 from keras.datasets import mnist
 
-from tmu.tsetlin_machine import TMCoalescedClassifier
+from tmu.models.classification.coalesced_classifier import TMCoalescedClassifier
 
 (X_train, Y_train), (X_test, Y_test) = mnist.load_data()
 
-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) 
+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 = TMCoalescedClassifier(20000, 50*100, 10.0, max_included_literals=5, weighted_clauses=True)
 
diff --git a/examples/experimental/classification/__init__.py b/examples/experimental/classification/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/examples/RegressionDemo.py b/examples/experimental/regression/RegressionDemo.py
similarity index 93%
rename from examples/RegressionDemo.py
rename to examples/experimental/regression/RegressionDemo.py
index 2ab0a150..637692f8 100644
--- a/examples/RegressionDemo.py
+++ b/examples/experimental/regression/RegressionDemo.py
@@ -1,4 +1,4 @@
-from tmu.tsetlin_machine import TMRegressor
+from tmu.models.regression.vanilla_regressor import TMRegressor
 from pyTsetlinMachine.tools import Binarizer
 import numpy as np
 from time import time
diff --git a/examples/experimental/regression/__init__.py b/examples/experimental/regression/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/examples/regression/__init__.py b/examples/regression/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/examples/requirements.txt b/examples/requirements.txt
new file mode 100644
index 00000000..480638e1
--- /dev/null
+++ b/examples/requirements.txt
@@ -0,0 +1,2 @@
+scikit-learn
+tensorflow
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 00000000..9e13d217
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,3 @@
+[build-system]
+requires = ["setuptools", "numpy", "wheel", "cffi>=1.0.0", "scipy"]
+build-backend = "setuptools.build_meta"
\ No newline at end of file
diff --git a/setup.py b/setup.py
index 5d2f710e..92323573 100644
--- a/setup.py
+++ b/setup.py
@@ -1,13 +1,12 @@
-import os
-import sys
 
+from pathlib import Path
 from setuptools import setup, find_packages
 
-os.chdir(os.path.dirname(sys.argv[0]) or ".")
+current_dir = Path(__file__).parent
 
 setup(
     name='tmu',
-    version='0.7.1',
+    version='0.7.2',
     url='https://github.com/cair/tmu/',
     author='Ole-Christoffer Granmo',
     author_email='ole.granmo@uia.no',
@@ -28,13 +27,11 @@
                      'autoencoder (https://arxiv.org/abs/2301.00709), literal budget (to be published), incremental '
                      'clause evaluation (to be published), and one-vs-one multi-class classifier (to be published). '
                      'TMU is written in Python with wrappers for C and CUDA-based clause evaluation and updating.',
-
-    setup_requires=["cffi>=1.0.0"],
     packages=find_packages(),
     cffi_modules=[
-        "./tmu/clause_bank_extension_build.py:ffibuilder",
-        "./tmu/tools_extension_build.py:ffibuilder",
-        "./tmu/weight_bank_extension_build.py:ffibuilder"
+        "tmu/lib/tmulib_extension_build.py:ffibuilder"
     ],
-    install_requires=["cffi>=1.0.0"],
-)
\ No newline at end of file
+    install_requires=[
+        "cffi>=1.0.0"
+    ]
+)
diff --git a/tmu/__init__.py b/tmu/__init__.py
index 8b137891..635180e5 100644
--- a/tmu/__init__.py
+++ b/tmu/__init__.py
@@ -1,1 +1,6 @@
 
+
+try:
+    import tmu.tmulib
+except ImportError:
+    raise ImportError("Could not import cffi compiled libraries. To fix this problem, run pip install -e .")
diff --git a/tmu/clause_bank.py b/tmu/clause_bank.py
index 93e4f95b..2ffdf1e4 100644
--- a/tmu/clause_bank.py
+++ b/tmu/clause_bank.py
@@ -21,173 +21,212 @@
 # This code implements the Convolutional Tsetlin Machine from paper arXiv:1905.09688
 # https://arxiv.org/abs/1905.09688
 
-from ._cb import ffi, lib
+from tmu.tmulib import ffi, lib
 
 import numpy as np
 
 import tmu.tools
 
-class ClauseBank():
-	def __init__(self, X, number_of_clauses, number_of_state_bits_ta, number_of_state_bits_ind, patch_dim, batch_size=100, incremental=True):
-		self.number_of_clauses = int(number_of_clauses)
-		self.number_of_state_bits_ta = int(number_of_state_bits_ta)
-		self.number_of_state_bits_ind = int(number_of_state_bits_ind)
-		self.patch_dim = patch_dim
-		self.batch_size = batch_size
-		self.incremental = incremental
 
-		if len(X.shape) == 2:
-			self.dim = (X.shape[1], 1, 1)
-		elif len(X.shape) == 3:
-			self.dim = (X.shape[1], X.shape[2], 1)
-		elif len(X.shape) == 4:
-			self.dim = (X.shape[1], X.shape[2], X.shape[3])
+class ClauseBank:
+    def __init__(self, X, number_of_clauses, number_of_state_bits_ta, number_of_state_bits_ind, patch_dim,
+                 batch_size=100, incremental=True):
+        self.number_of_clauses = int(number_of_clauses)
+        self.number_of_state_bits_ta = int(number_of_state_bits_ta)
+        self.number_of_state_bits_ind = int(number_of_state_bits_ind)
+        self.patch_dim = patch_dim
+        self.batch_size = batch_size
+        self.incremental = incremental
+
+        if len(X.shape) == 2:
+            self.dim = (X.shape[1], 1, 1)
+        elif len(X.shape) == 3:
+            self.dim = (X.shape[1], X.shape[2], 1)
+        elif len(X.shape) == 4:
+            self.dim = (X.shape[1], X.shape[2], X.shape[3])
 
-		if self.patch_dim == None:
-			self.patch_dim = (self.dim[0]*self.dim[1]*self.dim[2], 1)
-
-		self.number_of_features = int(self.patch_dim[0]*self.patch_dim[1]*self.dim[2] + (self.dim[0] - self.patch_dim[0]) + (self.dim[1] - self.patch_dim[1]))
-		self.number_of_literals = self.number_of_features*2
-		
-		self.number_of_patches = int((self.dim[0] - self.patch_dim[0] + 1)*(self.dim[1] - self.patch_dim[1] + 1))
-		self.number_of_ta_chunks = int((self.number_of_literals-1)/32 + 1)
-
-		self.clause_output = np.ascontiguousarray(np.empty((int(self.number_of_clauses)), dtype=np.uint32))
-		self.co_p = ffi.cast("unsigned int *", self.clause_output.ctypes.data)
-
-		self.clause_output_batch = np.ascontiguousarray(np.empty((int(self.number_of_clauses*batch_size)), dtype=np.uint32))
-		self.cob_p = ffi.cast("unsigned int *", self.clause_output_batch.ctypes.data)
-
-		self.clause_and_target = np.ascontiguousarray(np.zeros((int(self.number_of_clauses*self.number_of_ta_chunks)), dtype=np.uint32))
-		self.ct_p = ffi.cast("unsigned int *", self.clause_and_target.ctypes.data)
-
-		self.clause_output_patchwise = np.ascontiguousarray(np.empty((int(self.number_of_clauses*self.number_of_patches)), dtype=np.uint32))
-		self.cop_p = ffi.cast("unsigned int *", self.clause_output_patchwise.ctypes.data)
-
-		self.feedback_to_ta = np.ascontiguousarray(np.empty((self.number_of_ta_chunks), dtype=np.uint32))
-		self.ft_p = ffi.cast("unsigned int *", self.feedback_to_ta.ctypes.data)
-
-		self.output_one_patches = np.ascontiguousarray(np.empty(self.number_of_patches, dtype=np.uint32))
-		self.o1p_p = ffi.cast("unsigned int *", self.output_one_patches.ctypes.data)
-
-		self.literal_clause_count = np.ascontiguousarray(np.empty((int(self.number_of_literals)), dtype=np.uint32))
-		self.lcc_p = ffi.cast("unsigned int *", self.literal_clause_count.ctypes.data)
-
-		self.initialize_clauses()
-
-	def initialize_clauses(self):
-		self.clause_bank = np.empty((self.number_of_clauses, self.number_of_ta_chunks, self.number_of_state_bits_ta), dtype=np.uint32)
-		self.clause_bank[:,:,0:self.number_of_state_bits_ta-1] = np.uint32(~0)
-		self.clause_bank[:,:,self.number_of_state_bits_ta-1] = 0
-		self.clause_bank = np.ascontiguousarray(self.clause_bank.reshape((self.number_of_clauses * self.number_of_ta_chunks * self.number_of_state_bits_ta)))
-		self.cb_p = ffi.cast("unsigned int *", self.clause_bank.ctypes.data)
-
-		self.clause_bank_ind = np.empty((self.number_of_clauses, self.number_of_ta_chunks, self.number_of_state_bits_ind), dtype=np.uint32)
-		self.clause_bank_ind[:,:,:] = np.uint32(~0)
-		self.clause_bank_ind = np.ascontiguousarray(self.clause_bank_ind.reshape((self.number_of_clauses * self.number_of_ta_chunks * self.number_of_state_bits_ind)))
-		self.cbi_p = ffi.cast("unsigned int *", self.clause_bank_ind.ctypes.data)
-
-		self.incremental_clause_evaluation_initialized = False
-
-	def calculate_clause_outputs_predict(self, encoded_X, e):
-		xi_p = ffi.cast("unsigned int *", encoded_X[e,:].ctypes.data)
-
-		if not self.incremental:
-			lib.cb_calculate_clause_outputs_predict(self.cb_p, self.number_of_clauses, self.number_of_literals, self.number_of_state_bits_ta, self.number_of_patches, self.co_p, xi_p)
-			return self.clause_output
-
-		xi_p = ffi.cast("unsigned int *", encoded_X[e,:].ctypes.data)
-
-		if self.incremental_clause_evaluation_initialized == False:
-			self.literal_clause_map = np.ascontiguousarray(np.empty((int(self.number_of_literals*self.number_of_clauses)), dtype=np.uint32))
-			self.lcm_p = ffi.cast("unsigned int *", self.literal_clause_map.ctypes.data)
-
-			self.literal_clause_map_pos = np.ascontiguousarray(np.empty((int(self.number_of_literals)), dtype=np.uint32))
-			self.lcmp_p = ffi.cast("unsigned int *", self.literal_clause_map_pos.ctypes.data)
-
-			self.false_literals_per_clause = np.ascontiguousarray(np.empty((int(self.number_of_clauses*self.number_of_patches)), dtype=np.uint32))
-			self.flpc_p = ffi.cast("unsigned int *", self.false_literals_per_clause.ctypes.data)
-
-			self.previous_xi = np.ascontiguousarray(np.empty((int(self.number_of_ta_chunks)*int(self.number_of_patches)), dtype=np.uint32))
-			self.previous_xi_p = ffi.cast("unsigned int *", self.previous_xi.ctypes.data)
-
-			lib.cb_initialize_incremental_clause_calculation(self.cb_p, self.lcm_p, self.lcmp_p, self.flpc_p, self.number_of_clauses, self.number_of_literals, self.number_of_state_bits_ta, self.previous_xi_p)
-
-			self.incremental_clause_evaluation_initialized = True
-
-		if e % self.batch_size == 0:
-			lib.cb_calculate_clause_outputs_incremental_batch(self.lcm_p, self.lcmp_p, self.flpc_p, self.number_of_clauses, self.number_of_literals, self.number_of_patches, self.cob_p, self.previous_xi_p, xi_p, np.minimum(self.batch_size, encoded_X.shape[0]-e))
-
-		return self.clause_output_batch.reshape((self.batch_size, self.number_of_clauses))[e % self.batch_size,:]
-
-	def calculate_clause_outputs_update(self, literal_active, encoded_X, e):
-		xi_p = ffi.cast("unsigned int *", encoded_X[e,:].ctypes.data)
-		la_p = ffi.cast("unsigned int *", literal_active.ctypes.data)
-		lib.cb_calculate_clause_outputs_update(self.cb_p, self.number_of_clauses, self.number_of_literals, self.number_of_state_bits_ta, self.number_of_patches, self.co_p, la_p, xi_p)
-		return self.clause_output
-
-	def calculate_clause_outputs_patchwise(self, encoded_X, e):
-		xi_p = ffi.cast("unsigned int *", encoded_X[e,:].ctypes.data)
-		lib.cb_calculate_clause_outputs_patchwise(self.cb_p, self.number_of_clauses, self.number_of_literals, self.number_of_state_bits_ta, self.number_of_patches, self.cop_p, xi_p)
-		return self.clause_output_patchwise
-
-	def type_i_feedback(self, update_p, s, boost_true_positive_feedback, max_included_literals, clause_active, literal_active, encoded_X, e):
-		xi_p = ffi.cast("unsigned int *", encoded_X[e,:].ctypes.data)
-		ca_p = ffi.cast("unsigned int *", clause_active.ctypes.data)
-		la_p = ffi.cast("unsigned int *", literal_active.ctypes.data)
-		lib.cb_type_i_feedback(self.cb_p, self.ft_p, self.o1p_p, self.number_of_clauses, self.number_of_literals, self.number_of_state_bits_ta, self.number_of_patches, update_p, s, boost_true_positive_feedback, max_included_literals, ca_p, la_p, xi_p)
-		self.incremental_clause_evaluation_initialized = False
-
-	def type_ii_feedback(self, update_p, clause_active, literal_active, encoded_X, e):
-		xi_p = ffi.cast("unsigned int *", encoded_X[e,:].ctypes.data)
-		ca_p = ffi.cast("unsigned int *", clause_active.ctypes.data)
-		la_p = ffi.cast("unsigned int *", literal_active.ctypes.data)
-		lib.cb_type_ii_feedback(self.cb_p, self.o1p_p, self.number_of_clauses, self.number_of_literals, self.number_of_state_bits_ta, self.number_of_patches, update_p, ca_p, la_p, xi_p)
-		self.incremental_clause_evaluation_initialized = False
-
-	def type_iii_feedback(self, update_p, d, clause_active, literal_active, encoded_X, e, target):
-		xi_p = ffi.cast("unsigned int *", encoded_X[e,:].ctypes.data)
-		ca_p = ffi.cast("unsigned int *", clause_active.ctypes.data)
-		la_p = ffi.cast("unsigned int *", literal_active.ctypes.data)
-		lib.cb_type_iii_feedback(self.cb_p, self.cbi_p, self.ct_p, self.o1p_p, self.number_of_clauses, self.number_of_literals, self.number_of_state_bits_ta, self.number_of_state_bits_ind, self.number_of_patches, update_p, d, ca_p, la_p, xi_p, target)
-		self.incremental_clause_evaluation_initialized = False
-
-	def calculate_literal_clause_frequency(self, clause_active):
-		ca_p = ffi.cast("unsigned int *", clause_active.ctypes.data)
-		lib.cb_calculate_literal_frequency(self.cb_p, self.number_of_clauses, self.number_of_literals, self.number_of_state_bits_ta, ca_p, self.lcc_p)
-		return self.literal_clause_count
-	
-	def number_of_include_actions(self, clause):
-		return lib.cb_number_of_include_actions(self.cb_p, clause, self.number_of_literals, self.number_of_state_bits_ta)
-
-	def get_ta_action(self, clause, ta):
-		ta_chunk = ta // 32
-		chunk_pos = ta % 32
-		pos = int(clause * self.number_of_ta_chunks * self.number_of_state_bits_ta + ta_chunk * self.number_of_state_bits_ta + self.number_of_state_bits_ta-1)
-		return (self.clause_bank[pos] & (1 << chunk_pos)) > 0
-
-	def get_ta_state(self, clause, ta):
-		ta_chunk = ta // 32
-		chunk_pos = ta % 32
-		pos = int(clause * self.number_of_ta_chunks * self.number_of_state_bits_ta + ta_chunk * self.number_of_state_bits_ta)
-		state = 0
-		for b in range(self.number_of_state_bits_ta):
-			if self.clause_bank[pos + b] & (1 << chunk_pos) > 0:
-				state |= (1 << b)
-		return state
-
-	def set_ta_state(self, clause, ta, state):
-		ta_chunk = ta // 32
-		chunk_pos = ta % 32
-		pos = int(clause * self.number_of_ta_chunks * self.number_of_state_bits_ta + ta_chunk * self.number_of_state_bits_ta)
-		for b in range(self.number_of_state_bits_ta):
-			if state & (1 << b) > 0:
-				self.clause_bank[pos + b] |= (1 << chunk_pos)
-			else:
-				self.clause_bank[pos + b] &= ~(1 << chunk_pos)
-
-	def prepare_X(self, X):		
-		return tmu.tools.encode(X, X.shape[0], self.number_of_patches, self.number_of_ta_chunks, self.dim, self.patch_dim, 0)
-
-	def prepare_autoencoder_examples(self, X_csr, X_csc, active_output, accumulation):	
-		return tmu.tools.produce_autoencoder_examples(X_csr, X_csc, active_output, accumulation, self.number_of_ta_chunks)
+        if self.patch_dim == None:
+            self.patch_dim = (self.dim[0] * self.dim[1] * self.dim[2], 1)
+
+        self.number_of_features = int(
+            self.patch_dim[0] * self.patch_dim[1] * self.dim[2] + (self.dim[0] - self.patch_dim[0]) + (
+                        self.dim[1] - self.patch_dim[1]))
+        self.number_of_literals = self.number_of_features * 2
+
+        self.number_of_patches = int((self.dim[0] - self.patch_dim[0] + 1) * (self.dim[1] - self.patch_dim[1] + 1))
+        self.number_of_ta_chunks = int((self.number_of_literals - 1) / 32 + 1)
+
+        self.clause_output = np.ascontiguousarray(np.empty((int(self.number_of_clauses)), dtype=np.uint32))
+        self.co_p = ffi.cast("unsigned int *", self.clause_output.ctypes.data)
+
+        self.clause_output_batch = np.ascontiguousarray(
+            np.empty((int(self.number_of_clauses * batch_size)), dtype=np.uint32))
+        self.cob_p = ffi.cast("unsigned int *", self.clause_output_batch.ctypes.data)
+
+        self.clause_and_target = np.ascontiguousarray(
+            np.zeros((int(self.number_of_clauses * self.number_of_ta_chunks)), dtype=np.uint32))
+        self.ct_p = ffi.cast("unsigned int *", self.clause_and_target.ctypes.data)
+
+        self.clause_output_patchwise = np.ascontiguousarray(
+            np.empty((int(self.number_of_clauses * self.number_of_patches)), dtype=np.uint32))
+        self.cop_p = ffi.cast("unsigned int *", self.clause_output_patchwise.ctypes.data)
+
+        self.feedback_to_ta = np.ascontiguousarray(np.empty((self.number_of_ta_chunks), dtype=np.uint32))
+        self.ft_p = ffi.cast("unsigned int *", self.feedback_to_ta.ctypes.data)
+
+        self.output_one_patches = np.ascontiguousarray(np.empty(self.number_of_patches, dtype=np.uint32))
+        self.o1p_p = ffi.cast("unsigned int *", self.output_one_patches.ctypes.data)
+
+        self.literal_clause_count = np.ascontiguousarray(np.empty((int(self.number_of_literals)), dtype=np.uint32))
+        self.lcc_p = ffi.cast("unsigned int *", self.literal_clause_count.ctypes.data)
+
+        self.initialize_clauses()
+
+    def initialize_clauses(self):
+        self.clause_bank = np.empty((self.number_of_clauses, self.number_of_ta_chunks, self.number_of_state_bits_ta),
+                                    dtype=np.uint32)
+        self.clause_bank[:, :, 0:self.number_of_state_bits_ta - 1] = np.uint32(~0)
+        self.clause_bank[:, :, self.number_of_state_bits_ta - 1] = 0
+        self.clause_bank = np.ascontiguousarray(self.clause_bank.reshape(
+            (self.number_of_clauses * self.number_of_ta_chunks * self.number_of_state_bits_ta)))
+        self.cb_p = ffi.cast("unsigned int *", self.clause_bank.ctypes.data)
+
+        self.clause_bank_ind = np.empty(
+            (self.number_of_clauses, self.number_of_ta_chunks, self.number_of_state_bits_ind), dtype=np.uint32)
+        self.clause_bank_ind[:, :, :] = np.uint32(~0)
+        self.clause_bank_ind = np.ascontiguousarray(self.clause_bank_ind.reshape(
+            (self.number_of_clauses * self.number_of_ta_chunks * self.number_of_state_bits_ind)))
+        self.cbi_p = ffi.cast("unsigned int *", self.clause_bank_ind.ctypes.data)
+
+        self.incremental_clause_evaluation_initialized = False
+
+    def calculate_clause_outputs_predict(self, encoded_X, e):
+        xi_p = ffi.cast("unsigned int *", encoded_X[e, :].ctypes.data)
+
+        if not self.incremental:
+            lib.cb_calculate_clause_outputs_predict(self.cb_p, self.number_of_clauses, self.number_of_literals,
+                                                    self.number_of_state_bits_ta, self.number_of_patches, self.co_p,
+                                                    xi_p)
+            return self.clause_output
+
+        xi_p = ffi.cast("unsigned int *", encoded_X[e, :].ctypes.data)
+
+        if self.incremental_clause_evaluation_initialized == False:
+            self.literal_clause_map = np.ascontiguousarray(
+                np.empty((int(self.number_of_literals * self.number_of_clauses)), dtype=np.uint32))
+            self.lcm_p = ffi.cast("unsigned int *", self.literal_clause_map.ctypes.data)
+
+            self.literal_clause_map_pos = np.ascontiguousarray(
+                np.empty((int(self.number_of_literals)), dtype=np.uint32))
+            self.lcmp_p = ffi.cast("unsigned int *", self.literal_clause_map_pos.ctypes.data)
+
+            self.false_literals_per_clause = np.ascontiguousarray(
+                np.empty((int(self.number_of_clauses * self.number_of_patches)), dtype=np.uint32))
+            self.flpc_p = ffi.cast("unsigned int *", self.false_literals_per_clause.ctypes.data)
+
+            self.previous_xi = np.ascontiguousarray(
+                np.empty((int(self.number_of_ta_chunks) * int(self.number_of_patches)), dtype=np.uint32))
+            self.previous_xi_p = ffi.cast("unsigned int *", self.previous_xi.ctypes.data)
+
+            lib.cb_initialize_incremental_clause_calculation(self.cb_p, self.lcm_p, self.lcmp_p, self.flpc_p,
+                                                             self.number_of_clauses, self.number_of_literals,
+                                                             self.number_of_state_bits_ta, self.previous_xi_p)
+
+            self.incremental_clause_evaluation_initialized = True
+
+        if e % self.batch_size == 0:
+            lib.cb_calculate_clause_outputs_incremental_batch(self.lcm_p, self.lcmp_p, self.flpc_p,
+                                                              self.number_of_clauses, self.number_of_literals,
+                                                              self.number_of_patches, self.cob_p, self.previous_xi_p,
+                                                              xi_p, np.minimum(self.batch_size, encoded_X.shape[0] - e))
+
+        return self.clause_output_batch.reshape((self.batch_size, self.number_of_clauses))[e % self.batch_size, :]
+
+    def calculate_clause_outputs_update(self, literal_active, encoded_X, e):
+        xi_p = ffi.cast("unsigned int *", encoded_X[e, :].ctypes.data)
+        la_p = ffi.cast("unsigned int *", literal_active.ctypes.data)
+        lib.cb_calculate_clause_outputs_update(self.cb_p, self.number_of_clauses, self.number_of_literals,
+                                               self.number_of_state_bits_ta, self.number_of_patches, self.co_p, la_p,
+                                               xi_p)
+        return self.clause_output
+
+    def calculate_clause_outputs_patchwise(self, encoded_X, e):
+        xi_p = ffi.cast("unsigned int *", encoded_X[e, :].ctypes.data)
+        lib.cb_calculate_clause_outputs_patchwise(self.cb_p, self.number_of_clauses, self.number_of_literals,
+                                                  self.number_of_state_bits_ta, self.number_of_patches, self.cop_p,
+                                                  xi_p)
+        return self.clause_output_patchwise
+
+    def type_i_feedback(self, update_p, s, boost_true_positive_feedback, max_included_literals, clause_active,
+                        literal_active, encoded_X, e):
+        xi_p = ffi.cast("unsigned int *", encoded_X[e, :].ctypes.data)
+        ca_p = ffi.cast("unsigned int *", clause_active.ctypes.data)
+        la_p = ffi.cast("unsigned int *", literal_active.ctypes.data)
+        lib.cb_type_i_feedback(self.cb_p, self.ft_p, self.o1p_p, self.number_of_clauses, self.number_of_literals,
+                               self.number_of_state_bits_ta, self.number_of_patches, update_p, s,
+                               boost_true_positive_feedback, max_included_literals, ca_p, la_p, xi_p)
+        self.incremental_clause_evaluation_initialized = False
+
+    def type_ii_feedback(self, update_p, clause_active, literal_active, encoded_X, e):
+        xi_p = ffi.cast("unsigned int *", encoded_X[e, :].ctypes.data)
+        ca_p = ffi.cast("unsigned int *", clause_active.ctypes.data)
+        la_p = ffi.cast("unsigned int *", literal_active.ctypes.data)
+        lib.cb_type_ii_feedback(self.cb_p, self.o1p_p, self.number_of_clauses, self.number_of_literals,
+                                self.number_of_state_bits_ta, self.number_of_patches, update_p, ca_p, la_p, xi_p)
+        self.incremental_clause_evaluation_initialized = False
+
+    def type_iii_feedback(self, update_p, d, clause_active, literal_active, encoded_X, e, target):
+        xi_p = ffi.cast("unsigned int *", encoded_X[e, :].ctypes.data)
+        ca_p = ffi.cast("unsigned int *", clause_active.ctypes.data)
+        la_p = ffi.cast("unsigned int *", literal_active.ctypes.data)
+        lib.cb_type_iii_feedback(self.cb_p, self.cbi_p, self.ct_p, self.o1p_p, self.number_of_clauses,
+                                 self.number_of_literals, self.number_of_state_bits_ta, self.number_of_state_bits_ind,
+                                 self.number_of_patches, update_p, d, ca_p, la_p, xi_p, target)
+        self.incremental_clause_evaluation_initialized = False
+
+    def calculate_literal_clause_frequency(self, clause_active):
+        ca_p = ffi.cast("unsigned int *", clause_active.ctypes.data)
+        lib.cb_calculate_literal_frequency(self.cb_p, self.number_of_clauses, self.number_of_literals,
+                                           self.number_of_state_bits_ta, ca_p, self.lcc_p)
+        return self.literal_clause_count
+
+    def number_of_include_actions(self, clause):
+        return lib.cb_number_of_include_actions(self.cb_p, clause, self.number_of_literals,
+                                                self.number_of_state_bits_ta)
+
+    def get_ta_action(self, clause, ta):
+        ta_chunk = ta // 32
+        chunk_pos = ta % 32
+        pos = int(
+            clause * self.number_of_ta_chunks * self.number_of_state_bits_ta + ta_chunk * self.number_of_state_bits_ta + self.number_of_state_bits_ta - 1)
+        return (self.clause_bank[pos] & (1 << chunk_pos)) > 0
+
+    def get_ta_state(self, clause, ta):
+        ta_chunk = ta // 32
+        chunk_pos = ta % 32
+        pos = int(
+            clause * self.number_of_ta_chunks * self.number_of_state_bits_ta + ta_chunk * self.number_of_state_bits_ta)
+        state = 0
+        for b in range(self.number_of_state_bits_ta):
+            if self.clause_bank[pos + b] & (1 << chunk_pos) > 0:
+                state |= (1 << b)
+        return state
+
+    def set_ta_state(self, clause, ta, state):
+        ta_chunk = ta // 32
+        chunk_pos = ta % 32
+        pos = int(
+            clause * self.number_of_ta_chunks * self.number_of_state_bits_ta + ta_chunk * self.number_of_state_bits_ta)
+        for b in range(self.number_of_state_bits_ta):
+            if state & (1 << b) > 0:
+                self.clause_bank[pos + b] |= (1 << chunk_pos)
+            else:
+                self.clause_bank[pos + b] &= ~(1 << chunk_pos)
+
+    def prepare_X(self, X):
+        return tmu.tools.encode(X, X.shape[0], self.number_of_patches, self.number_of_ta_chunks, self.dim,
+                                self.patch_dim, 0)
+
+    def prepare_autoencoder_examples(self, X_csr, X_csc, active_output, accumulation):
+        return tmu.tools.produce_autoencoder_examples(X_csr, X_csc, active_output, accumulation,
+                                                      self.number_of_ta_chunks)
diff --git a/tmu/clause_bank_cuda.py b/tmu/clause_bank_cuda.py
index f1fdd018..38d6f467 100644
--- a/tmu/clause_bank_cuda.py
+++ b/tmu/clause_bank_cuda.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2022 Ole-Christoffer Granmo
+# Copyright (c) 2023 Ole-Christoffer Granmo
 
 # Permission is hereby granted, free of charge, to any person obtaining a copy
 # of this software and associated documentation files (the "Software"), to deal
@@ -21,7 +21,7 @@
 # This code implements the Convolutional Tsetlin Machine from paper arXiv:1905.09688
 # https://arxiv.org/abs/1905.09688
 
-from ._cb import ffi, lib
+from tmu.tmulib import ffi, lib
 
 import numpy as np
 
diff --git a/tmu/clause_bank_extension_build.py b/tmu/clause_bank_extension_build.py
deleted file mode 100644
index 51daabf0..00000000
--- a/tmu/clause_bank_extension_build.py
+++ /dev/null
@@ -1,42 +0,0 @@
-# Copyright (c) 2023 Ole-Christoffer Granmo
-
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-# SOFTWARE.
-
-# This code implements the Convolutional Tsetlin Machine from paper arXiv:1905.09688
-# https://arxiv.org/abs/1905.09688
-
-from cffi import FFI
-import pathlib
-current_dir = pathlib.Path(__file__).parent
-
-ffibuilder  = FFI()
-
-with current_dir.joinpath("ClauseBank.h").open("r") as f:
-    ffibuilder .cdef(f.read())
-
-with current_dir.joinpath("ClauseBank.c").open("r") as f:
-    ffibuilder .set_source(
-        "tmu._cb",
-        f.read(),
-        include_dirs=[current_dir.absolute()],
-        source_extension=".c",
-)
-
-if __name__ == "__main__":
-    ffibuilder .compile(verbose=True)
diff --git a/tmu/clause_bank_numba.py b/tmu/clause_bank_numba.py
deleted file mode 100644
index 8b137891..00000000
--- a/tmu/clause_bank_numba.py
+++ /dev/null
@@ -1,1 +0,0 @@
-
diff --git a/tmu/clause_bank_parallel.py b/tmu/clause_bank_parallel.py
deleted file mode 100644
index 8b137891..00000000
--- a/tmu/clause_bank_parallel.py
+++ /dev/null
@@ -1,1 +0,0 @@
-
diff --git a/tmu/datasets.py b/tmu/datasets.py
new file mode 100644
index 00000000..c6855011
--- /dev/null
+++ b/tmu/datasets.py
@@ -0,0 +1,43 @@
+import abc
+from typing import Dict
+
+import numpy as np
+
+
+class TMUDataset:
+
+    def __init__(self):
+        pass
+
+    @abc.abstractmethod
+    def booleanizer(self, name, dataset):
+        raise NotImplementedError("You should override def threshold()")
+
+    @abc.abstractmethod
+    def retrieve_dataset(self) -> Dict[str, np.ndarray]:
+        raise NotImplementedError("You should override def retrieve_dataset()")
+
+    def get(self):
+        return {k: self.booleanizer(k, v) for k, v in self.retrieve_dataset().items()}
+
+    def get_list(self):
+        return list(self.get().values())
+
+
+class MNIST(TMUDataset):
+    def retrieve_dataset(self) -> Dict[str, np.ndarray]:
+        from keras.datasets import mnist
+        (X_train, Y_train), (X_test, Y_test) = mnist.load_data()
+        return dict(
+            x_train=X_train,
+            y_train=Y_train,
+            x_test=X_test,
+            y_test=Y_test
+        )
+
+    def booleanizer(self, name, dataset):
+        if name.startswith("y"):
+            return dataset
+
+        return np.where(dataset.reshape((dataset.shape[0], 28*28)) > 75, 1, 0)
+
diff --git a/tmu/ClauseBank.h b/tmu/lib/include/ClauseBank.h
similarity index 100%
rename from tmu/ClauseBank.h
rename to tmu/lib/include/ClauseBank.h
diff --git a/tmu/Tools.h b/tmu/lib/include/Tools.h
similarity index 100%
rename from tmu/Tools.h
rename to tmu/lib/include/Tools.h
diff --git a/tmu/WeightBank.h b/tmu/lib/include/WeightBank.h
similarity index 100%
rename from tmu/WeightBank.h
rename to tmu/lib/include/WeightBank.h
diff --git a/tmu/fast_rand.h b/tmu/lib/include/fast_rand.h
similarity index 100%
rename from tmu/fast_rand.h
rename to tmu/lib/include/fast_rand.h
diff --git a/tmu/ClauseBank.c b/tmu/lib/src/ClauseBank.c
similarity index 100%
rename from tmu/ClauseBank.c
rename to tmu/lib/src/ClauseBank.c
diff --git a/tmu/Tools.c b/tmu/lib/src/Tools.c
similarity index 100%
rename from tmu/Tools.c
rename to tmu/lib/src/Tools.c
diff --git a/tmu/WeightBank.c b/tmu/lib/src/WeightBank.c
similarity index 100%
rename from tmu/WeightBank.c
rename to tmu/lib/src/WeightBank.c
diff --git a/tmu/tools_extension_build.py b/tmu/lib/tmulib_extension_build.py
similarity index 70%
rename from tmu/tools_extension_build.py
rename to tmu/lib/tmulib_extension_build.py
index 070b66eb..504fd207 100644
--- a/tmu/tools_extension_build.py
+++ b/tmu/lib/tmulib_extension_build.py
@@ -23,20 +23,25 @@
 
 from cffi import FFI
 import pathlib
+
 current_dir = pathlib.Path(__file__).parent
+source_dir = current_dir.joinpath("src")
+header_dir = current_dir.joinpath("include")
 
-ffibuilder = FFI()
+HEADERS = ["ClauseBank.h", "Tools.h", "WeightBank.h"]
+SOURCES = ["ClauseBank.c", "Tools.c", "WeightBank.c"]
 
-with current_dir.joinpath("Tools.h").open("r") as f:
-    ffibuilder .cdef(f.read())
+header_content = '\n'.join([header_dir.joinpath(x).open("r").read() for x in HEADERS])
+source_content = '\n'.join([source_dir.joinpath(x).open("r").read() for x in SOURCES])
 
-with current_dir.joinpath("Tools.c").open("r") as f:
-    ffibuilder .set_source(
-        "tmu._tools",
-        f.read(),
-        include_dirs=[current_dir.absolute()],
-        source_extension=".c",
-    )
+ffibuilder = FFI()
+ffibuilder.cdef(header_content)
+ffibuilder.set_source(
+    "tmu.tmulib",
+    source_content,
+    include_dirs=[header_dir.absolute()],
+    extra_compile_args=["-O3"]
+)
 
 if __name__ == "__main__":
-    ffibuilder .compile(verbose=True)
+    ffibuilder.compile(verbose=True)
diff --git a/tmu/models/__init__.py b/tmu/models/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tmu/models/autoencoder/__init__.py b/tmu/models/autoencoder/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tmu/models/autoencoder/autoencoder.py b/tmu/models/autoencoder/autoencoder.py
new file mode 100644
index 00000000..14bd4091
--- /dev/null
+++ b/tmu/models/autoencoder/autoencoder.py
@@ -0,0 +1,275 @@
+# Copyright (c) 2023 Ole-Christoffer Granmo
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+from tmu.clause_bank import ClauseBank
+from tmu.models.base import TMBasis
+from tmu.weight_bank import WeightBank
+import numpy as np
+from scipy.sparse import csr_matrix, csc_matrix
+
+
+class TMAutoEncoder(TMBasis):
+    def __init__(self, number_of_clauses, T, s, output_active, accumulation=1, type_i_ii_ratio=1.0,
+                 type_iii_feedback=False, focused_negative_sampling=False, output_balancing=False, d=200.0,
+                 platform='CPU', patch_dim=None, feature_negation=True, boost_true_positive_feedback=1,
+                 max_included_literals=None, number_of_state_bits_ta=8, number_of_state_bits_ind=8,
+                 weighted_clauses=False, clause_drop_p=0.0, literal_drop_p=0.0):
+        self.output_active = output_active
+        self.accumulation = accumulation
+        super().__init__(number_of_clauses, T, s, type_i_ii_ratio=type_i_ii_ratio, type_iii_feedback=type_iii_feedback,
+                         focused_negative_sampling=focused_negative_sampling, output_balancing=output_balancing, d=d,
+                         platform=platform, patch_dim=patch_dim, feature_negation=feature_negation,
+                         boost_true_positive_feedback=boost_true_positive_feedback,
+                         max_included_literals=max_included_literals, number_of_state_bits_ta=number_of_state_bits_ta,
+                         number_of_state_bits_ind=number_of_state_bits_ind, weighted_clauses=weighted_clauses,
+                         clause_drop_p=clause_drop_p, literal_drop_p=literal_drop_p)
+
+    def initialize(self, X):
+        self.number_of_classes = self.output_active.shape[0]
+        if self.platform == 'CPU':
+            self.clause_bank = ClauseBank(X, self.number_of_clauses, self.number_of_state_bits_ta,
+                                          self.number_of_state_bits_ind, self.patch_dim)
+        elif self.platform == 'CUDA':
+            from tmu.clause_bank_cuda import ClauseBankCUDA
+            self.clause_bank = ClauseBankCUDA(X, self.number_of_clauses, self.number_of_state_bits_ta, self.patch_dim)
+        else:
+            raise RuntimeError(f"Unknown platform of type: {self.platform}")
+
+        self.weight_banks = []
+        for i in range(self.number_of_classes):
+            self.weight_banks.append(
+                WeightBank(np.random.choice([-1, 1], size=self.number_of_clauses).astype(np.int32)))
+
+        if self.max_included_literals == None:
+            self.max_included_literals = self.clause_bank.number_of_literals
+
+    def update(self, target_output, target_value, encoded_X, clause_active, literal_active):
+        all_literal_active = (np.zeros(self.clause_bank.number_of_ta_chunks, dtype=np.uint32) | ~0).astype(np.uint32)
+        clause_outputs = self.clause_bank.calculate_clause_outputs_update(all_literal_active, encoded_X, 0)
+
+        class_sum = np.dot(clause_active * self.weight_banks[target_output].get_weights(), clause_outputs).astype(
+            np.int32)
+        class_sum = np.clip(class_sum, -self.T, self.T)
+
+        type_iii_feedback_selection = np.random.choice(2)
+
+        if target_value == 1:
+            update_p = (self.T - class_sum) / (2 * self.T)
+
+            self.clause_bank.type_i_feedback(update_p * self.type_i_p, self.s, self.boost_true_positive_feedback,
+                                             self.max_included_literals,
+                                             clause_active * (self.weight_banks[target_output].get_weights() >= 0),
+                                             literal_active, encoded_X, 0)
+            self.clause_bank.type_ii_feedback(update_p * self.type_ii_p,
+                                              clause_active * (self.weight_banks[target_output].get_weights() < 0),
+                                              literal_active, encoded_X, 0)
+            self.weight_banks[target_output].increment(clause_outputs, update_p, clause_active, True)
+            if self.type_iii_feedback and type_iii_feedback_selection == 0:
+                self.clause_bank.type_iii_feedback(update_p, self.d, clause_active * (
+                        self.weight_banks[target_output].get_weights() >= 0), literal_active, encoded_X, 0, 1)
+                self.clause_bank.type_iii_feedback(update_p, self.d,
+                                                   clause_active * (self.weight_banks[target_output].get_weights() < 0),
+                                                   literal_active, encoded_X, 0, 0)
+        else:
+            update_p = (self.T + class_sum) / (2 * self.T)
+
+            self.clause_bank.type_i_feedback(update_p * self.type_i_p, self.s, self.boost_true_positive_feedback,
+                                             self.max_included_literals,
+                                             clause_active * (self.weight_banks[target_output].get_weights() < 0),
+                                             literal_active, encoded_X, 0)
+            self.clause_bank.type_ii_feedback(update_p * self.type_ii_p,
+                                              clause_active * (self.weight_banks[target_output].get_weights() >= 0),
+                                              literal_active, encoded_X, 0)
+            self.weight_banks[target_output].decrement(clause_outputs, update_p, clause_active, True)
+            if self.type_iii_feedback and type_iii_feedback_selection == 1:
+                self.clause_bank.type_iii_feedback(update_p, self.d,
+                                                   clause_active * (self.weight_banks[target_output].get_weights() < 0),
+                                                   literal_active, encoded_X, 0, 1)
+                self.clause_bank.type_iii_feedback(update_p, self.d, clause_active * (
+                        self.weight_banks[target_output].get_weights() >= 0), literal_active, encoded_X, 0, 0)
+        return
+
+    def activate_clauses(self):
+        # Clauses are dropped based on their weights
+        clause_active = np.ones(self.number_of_clauses, dtype=np.uint32)
+        deactivate = np.random.choice(self.number_of_clauses, size=int(self.number_of_clauses * self.clause_drop_p))
+        for d in range(deactivate.shape[0]):
+            clause_active[deactivate[d]] = 0
+
+        return clause_active
+
+    def activate_literals(self):
+        # Literals are dropped based on their frequency
+        literal_active = (np.zeros(self.clause_bank.number_of_ta_chunks, dtype=np.uint32) | ~0).astype(np.uint32)
+        literal_clause_frequency = self.literal_clause_frequency()
+        deactivate = np.random.choice(self.clause_bank.number_of_literals,
+                                      size=int(self.clause_bank.number_of_literals * self.literal_drop_p))
+        for d in range(deactivate.shape[0]):
+            ta_chunk = deactivate[d] // 32
+            chunk_pos = deactivate[d] % 32
+            literal_active[ta_chunk] &= (~(1 << chunk_pos))
+
+        if not self.feature_negation:
+            for k in range(self.clause_bank.number_of_literals // 2, self.clause_bank.number_of_literals):
+                ta_chunk = k // 32
+                chunk_pos = k % 32
+                literal_active[ta_chunk] &= (~(1 << chunk_pos))
+        literal_active = literal_active.astype(np.uint32)
+
+        return (literal_active)
+
+    def fit(self, X, number_of_examples=2000, shuffle=True):
+        if self.initialized == False:
+            self.initialize(X)
+            self.initialized = True
+
+        X_csr = csr_matrix(X.reshape(X.shape[0], -1))
+        X_csc = csc_matrix(X.reshape(X.shape[0], -1)).sorted_indices()
+
+        clause_active = self.activate_clauses()
+        literal_active = self.activate_literals()
+
+        class_index = np.arange(self.number_of_classes, dtype=np.uint32)
+        for e in range(number_of_examples):
+            np.random.shuffle(class_index)
+
+            average_absolute_weights = np.zeros(self.number_of_clauses, dtype=np.float32)
+            for i in class_index:
+                average_absolute_weights += np.absolute(self.weight_banks[i].get_weights())
+            average_absolute_weights /= self.number_of_classes
+            update_clause = np.random.random(self.number_of_clauses) <= (
+                    self.T - np.clip(average_absolute_weights, 0, self.T)) / self.T
+
+            Xu, Yu = self.clause_bank.prepare_autoencoder_examples(X_csr, X_csc, self.output_active, self.accumulation)
+            for i in class_index:
+                (target, encoded_X) = Yu[i], Xu[i].reshape((1, -1))
+
+                ta_chunk = self.output_active[i] // 32
+                chunk_pos = self.output_active[i] % 32
+                copy_literal_active_ta_chunk = literal_active[ta_chunk]
+                literal_active[ta_chunk] &= ~(1 << chunk_pos)
+
+                self.update(i, target, encoded_X, update_clause * clause_active, literal_active)
+                literal_active[ta_chunk] = copy_literal_active_ta_chunk
+        return
+
+    def predict(self, X):
+        X_csr = csr_matrix(X.reshape(X.shape[0], -1))
+        Y = np.ascontiguousarray(np.zeros((self.number_of_classes, X.shape[0]), dtype=np.uint32))
+
+        for e in range(X.shape[0]):
+            encoded_X = self.clause_bank.prepare_X(X_csr[e, :].toarray())
+
+            clause_outputs = self.clause_bank.calculate_clause_outputs_predict(encoded_X, 0)[0]
+            for i in range(self.number_of_classes):
+                class_sum = np.dot(self.weight_banks[i].get_weights(), clause_outputs).astype(np.int32)
+                Y[i, e] = (class_sum >= 0)
+        return Y
+
+    def literal_importance(self, the_class, negated_features=False, negative_polarity=False):
+        literal_frequency = np.zeros(self.clause_bank.number_of_literals, dtype=np.uint32)
+        if negated_features:
+            if negative_polarity:
+                literal_frequency[
+                self.clause_bank.number_of_literals // 2:] += self.clause_bank.calculate_literal_clause_frequency(
+                    self.weight_banks[the_class].get_weights() < 0)[self.clause_bank.number_of_literals // 2:]
+            else:
+                literal_frequency[
+                self.clause_bank.number_of_literals // 2:] += self.clause_bank.calculate_literal_clause_frequency(
+                    self.weight_banks[the_class].get_weights() >= 0)[self.clause_bank.number_of_literals // 2:]
+        else:
+            if negative_polarity:
+                literal_frequency[
+                :self.clause_bank.number_of_literals // 2] += self.clause_bank.calculate_literal_clause_frequency(
+                    self.weight_banks[the_class].get_weights() < 0)[:self.clause_bank.number_of_literals // 2]
+            else:
+                literal_frequency[
+                :self.clause_bank.number_of_literals // 2] += self.clause_bank.calculate_literal_clause_frequency(
+                    self.weight_banks[the_class].get_weights() >= 0)[:self.clause_bank.number_of_literals // 2]
+
+        return literal_frequency
+
+    def clause_precision(self, the_class, positive_polarity, X, number_of_examples=2000):
+        X_csr = csr_matrix(X.reshape(X.shape[0], -1))
+        X_csc = csc_matrix(X.reshape(X.shape[0], -1)).sorted_indices()
+
+        true_positive = np.zeros(self.number_of_clauses, dtype=np.uint32)
+        false_positive = np.zeros(self.number_of_clauses, dtype=np.uint32)
+
+        weights = self.weight_banks[the_class].get_weights()
+
+        clause_active = self.activate_clauses()
+        literal_active = self.activate_literals()
+
+        for e in range(number_of_examples):
+            Xu, Yu = self.clause_bank.prepare_autoencoder_examples(X_csr, X_csc,
+                                                                   np.array([self.output_active[the_class]],
+                                                                            dtype=np.uint32), self.accumulation)
+            (target, encoded_X) = Yu[0], Xu[0].reshape((1, -1))
+
+            clause_outputs = self.clause_bank.calculate_clause_outputs_predict(encoded_X, 0)
+
+            if positive_polarity:
+                if target == 1:
+                    true_positive += (weights >= 0) * clause_outputs
+                else:
+                    false_positive += (weights >= 0) * clause_outputs
+            else:
+                if target == 0:
+                    true_positive += (weights < 0) * clause_outputs
+                else:
+                    false_positive += (weights < 0) * clause_outputs
+
+        return 1.0 * true_positive / (true_positive + false_positive)
+
+    def clause_recall(self, the_class, positive_polarity, X, number_of_examples=2000):
+        X_csr = csr_matrix(X.reshape(X.shape[0], -1))
+        X_csc = csc_matrix(X.reshape(X.shape[0], -1)).sorted_indices()
+
+        true_positive = np.zeros(self.number_of_clauses, dtype=np.uint32)
+        false_negative = np.zeros(self.number_of_clauses, dtype=np.uint32)
+
+        weights = self.weight_banks[the_class].get_weights()
+
+        for e in range(number_of_examples):
+            Xu, Yu = self.clause_bank.prepare_autoencoder_examples(X_csr, X_csc,
+                                                                   np.array([self.output_active[the_class]],
+                                                                            dtype=np.uint32), self.accumulation)
+            (target, encoded_X) = Yu[0], Xu[0].reshape((1, -1))
+
+            clause_outputs = self.clause_bank.calculate_clause_outputs_predict(encoded_X, 0)
+
+            if positive_polarity:
+                if target == 1:
+                    true_positive += (weights >= 0) * clause_outputs
+                    false_negative += (weights >= 0) * (1 - clause_outputs)
+            else:
+                if target == 0:
+                    true_positive += (weights < 0) * clause_outputs
+                    false_negative += (weights < 0) * (1 - clause_outputs)
+
+        return true_positive / (true_positive + false_negative)
+
+    def get_weight(self, the_class, clause):
+        return self.weight_banks[the_class].get_weights()[clause]
+
+    def get_weights(self, the_class):
+        return self.weight_banks[the_class].get_weights()
+
+    def set_weight(self, the_class, clause, weight):
+        self.weight_banks[the_class].get_weights()[clause] = weight
diff --git a/tmu/models/base.py b/tmu/models/base.py
new file mode 100644
index 00000000..ceae122b
--- /dev/null
+++ b/tmu/models/base.py
@@ -0,0 +1,102 @@
+# Copyright (c) 2023 Ole-Christoffer Granmo
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+import numpy as np
+from scipy.sparse import csr_matrix
+import tmu.tools
+
+
+class TMBasis:
+    def __init__(self, number_of_clauses, T, s, confidence_driven_updating=False, type_i_ii_ratio=1.0,
+                 type_iii_feedback=False, focused_negative_sampling=False, output_balancing=False, d=200.0,
+                 platform='CPU', patch_dim=None, feature_negation=True, boost_true_positive_feedback=1,
+                 max_included_literals=None, number_of_state_bits_ta=8, number_of_state_bits_ind=8,
+                 weighted_clauses=False, clause_drop_p=0.0, literal_drop_p=0.0, batch_size=100, incremental=True):
+        self.number_of_clauses = number_of_clauses
+        self.number_of_state_bits_ta = number_of_state_bits_ta
+        self.number_of_state_bits_ind = number_of_state_bits_ind
+        self.T = int(T)
+        self.s = s
+
+        self.confidence_driven_updating = confidence_driven_updating
+
+        if type_i_ii_ratio >= 1.0:
+            self.type_i_p = 1.0
+            self.type_ii_p = 1.0 / type_i_ii_ratio
+        else:
+            self.type_i_p = type_i_ii_ratio
+            self.type_ii_p = 1.0
+
+        self.type_iii_feedback = type_iii_feedback
+        self.focused_negative_sampling = focused_negative_sampling
+        self.output_balancing = output_balancing
+        self.d = d
+        self.platform = platform
+        self.patch_dim = patch_dim
+        self.feature_negation = feature_negation
+        self.boost_true_positive_feedback = boost_true_positive_feedback
+        self.max_included_literals = max_included_literals
+        self.weighted_clauses = weighted_clauses
+
+        self.clause_drop_p = clause_drop_p
+        self.literal_drop_p = literal_drop_p
+
+        self.batch_size = batch_size
+
+        self.incremental = incremental
+
+        self.X_train = np.zeros(0, dtype=np.uint32)
+        self.X_test = np.zeros(0, dtype=np.uint32)
+
+        self.initialized = False
+
+    def clause_co_occurrence(self, X, percentage=False):
+        clause_outputs = csr_matrix(self.transform(X))
+        if percentage:
+            return clause_outputs.transpose().dot(clause_outputs).multiply(1.0 / clause_outputs.sum(axis=0))
+        else:
+            return clause_outputs.transpose().dot(clause_outputs)
+
+    def transform(self, X):
+        encoded_X = self.clause_bank.prepare_X(X)
+        transformed_X = np.empty((X.shape[0], self.number_of_clauses), dtype=np.uint32)
+        for e in range(X.shape[0]):
+            transformed_X[e, :] = self.clause_bank.calculate_clause_outputs_predict(encoded_X, e)
+        return transformed_X
+
+    def transform_patchwise(self, X):
+        encoded_X = tmu.tools.encode(X, X.shape[0], self.number_of_patches, self.clause_bank.number_of_ta_chunks,
+                                     self.dim, self.patch_dim, 0)
+        transformed_X = np.empty((X.shape[0], self.number_of_clauses * self.number_of_patches), dtype=np.uint32)
+        for e in range(X.shape[0]):
+            transformed_X[e, :] = self.clause_bank.calculate_clause_outputs_patchwise(encoded_X, e)
+        return transformed_X.reshape((X.shape[0], self.number_of_clauses, self.number_of_patches))
+
+    def literal_clause_frequency(self):
+        clause_active = np.ones(self.number_of_clauses, dtype=np.uint32)
+        return self.clause_bank.calculate_literal_clause_frequency(clause_active)
+
+    def get_ta_action(self, clause, ta):
+        return self.clause_bank.get_ta_action(clause, ta)
+
+    def get_ta_state(self, clause, ta):
+        return self.clause_bank.get_ta_state(clause, ta)
+
+    def set_ta_state(self, clause, ta, state):
+        return self.clause_bank.set_ta_state(clause, ta, state)
diff --git a/tmu/models/classification/__init__.py b/tmu/models/classification/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tmu/models/classification/coalesced_classifier.py b/tmu/models/classification/coalesced_classifier.py
new file mode 100644
index 00000000..e08efbc4
--- /dev/null
+++ b/tmu/models/classification/coalesced_classifier.py
@@ -0,0 +1,270 @@
+# Copyright (c) 2023 Ole-Christoffer Granmo
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+from tmu.clause_bank import ClauseBank
+from tmu.models.base import TMBasis
+from tmu.weight_bank import WeightBank
+import numpy as np
+
+
+class TMCoalescedClassifier(TMBasis):
+    def __init__(self, number_of_clauses, T, s, type_i_ii_ratio=1.0, type_iii_feedback=False,
+                 focused_negative_sampling=False, output_balancing=False, d=200.0, platform='CPU', patch_dim=None,
+                 feature_negation=True, boost_true_positive_feedback=1, max_included_literals=None,
+                 number_of_state_bits_ta=8, number_of_state_bits_ind=8, weighted_clauses=False, clause_drop_p=0.0,
+                 literal_drop_p=0.0):
+        super().__init__(number_of_clauses, T, s, type_i_ii_ratio=type_i_ii_ratio, type_iii_feedback=type_iii_feedback,
+                         focused_negative_sampling=focused_negative_sampling, output_balancing=output_balancing, d=d,
+                         platform=platform, patch_dim=patch_dim, feature_negation=feature_negation,
+                         boost_true_positive_feedback=boost_true_positive_feedback,
+                         max_included_literals=max_included_literals, number_of_state_bits_ta=number_of_state_bits_ta,
+                         number_of_state_bits_ind=number_of_state_bits_ind, weighted_clauses=weighted_clauses,
+                         clause_drop_p=clause_drop_p, literal_drop_p=literal_drop_p)
+
+    def initialize(self, X, Y):
+        self.number_of_classes = int(np.max(Y) + 1)
+
+        if self.platform == 'CPU':
+            self.clause_bank = ClauseBank(X, self.number_of_clauses, self.number_of_state_bits_ta,
+                                          self.number_of_state_bits_ind, self.patch_dim)
+        elif self.platform == 'CUDA':
+            from tmu.clause_bank_cuda import ClauseBankCUDA
+            self.clause_bank = ClauseBankCUDA(X, self.number_of_clauses, self.number_of_state_bits_ta, self.patch_dim)
+        else:
+            raise RuntimeError(f"Unknown platform of type: {self.platform}")
+
+        self.weight_banks = []
+        for i in range(self.number_of_classes):
+            self.weight_banks.append(
+                WeightBank(np.random.choice([-1, 1], size=self.number_of_clauses).astype(np.int32)))
+
+        if self.max_included_literals == None:
+            self.max_included_literals = self.clause_bank.number_of_literals
+
+    def update(self, target, e):
+        clause_outputs = self.clause_bank.calculate_clause_outputs_update(self.literal_active, self.encoded_X_train, e)
+
+        class_sum = np.dot(self.clause_active * self.weight_banks[target].get_weights(), clause_outputs).astype(
+            np.int32)
+        class_sum = np.clip(class_sum, -self.T, self.T)
+        update_p = (self.T - class_sum) / (2 * self.T)
+
+        type_iii_feedback_selection = np.random.choice(2)
+
+        self.clause_bank.type_i_feedback(update_p * self.type_i_p, self.s, self.boost_true_positive_feedback,
+                                         self.max_included_literals,
+                                         self.clause_active * (self.weight_banks[target].get_weights() >= 0),
+                                         self.literal_active, self.encoded_X_train, e)
+        self.clause_bank.type_ii_feedback(update_p * self.type_ii_p,
+                                          self.clause_active * (self.weight_banks[target].get_weights() < 0),
+                                          self.literal_active, self.encoded_X_train, e)
+        self.weight_banks[target].increment(clause_outputs, update_p, self.clause_active, True)
+        if self.type_iii_feedback and type_iii_feedback_selection == 0:
+            self.clause_bank.type_iii_feedback(update_p, self.d,
+                                               self.clause_active * (self.weight_banks[target].get_weights() >= 0),
+                                               self.literal_active, self.encoded_X_train, e, 1)
+            self.clause_bank.type_iii_feedback(update_p, self.d,
+                                               self.clause_active * (self.weight_banks[target].get_weights() < 0),
+                                               self.literal_active, self.encoded_X_train, e, 0)
+
+        for i in range(self.number_of_classes):
+            if i == target:
+                self.update_ps[i] = 0.0
+            else:
+                self.update_ps[i] = np.dot(self.clause_active * self.weight_banks[i].get_weights(),
+                                           clause_outputs).astype(np.int32)
+                self.update_ps[i] = np.clip(self.update_ps[i], -self.T, self.T)
+                self.update_ps[i] = 1.0 * (self.T + self.update_ps[i]) / (2 * self.T)
+
+        if self.update_ps.sum() == 0:
+            return
+
+        if self.focused_negative_sampling:
+            not_target = np.random.choice(self.number_of_classes, p=self.update_ps / self.update_ps.sum())
+            update_p = self.update_ps[not_target]
+        else:
+            not_target = np.random.randint(self.number_of_classes)
+            while not_target == target:
+                not_target = np.random.randint(self.number_of_classes)
+            update_p = self.update_ps[not_target]
+
+        self.clause_bank.type_i_feedback(update_p * self.type_i_p, self.s, self.boost_true_positive_feedback,
+                                         self.max_included_literals,
+                                         self.clause_active * (self.weight_banks[not_target].get_weights() < 0),
+                                         self.literal_active, self.encoded_X_train, e)
+        self.clause_bank.type_ii_feedback(update_p * self.type_ii_p,
+                                          self.clause_active * (self.weight_banks[not_target].get_weights() >= 0),
+                                          self.literal_active, self.encoded_X_train, e)
+        if self.type_iii_feedback and type_iii_feedback_selection == 1:
+            self.clause_bank.type_iii_feedback(update_p, self.d,
+                                               self.clause_active * (self.weight_banks[not_target].get_weights() < 0),
+                                               self.literal_active, self.encoded_X_train, e, 1)
+            self.clause_bank.type_iii_feedback(update_p, self.d,
+                                               self.clause_active * (self.weight_banks[not_target].get_weights() >= 0),
+                                               self.literal_active, self.encoded_X_train, e, 0)
+
+        self.weight_banks[not_target].decrement(clause_outputs, update_p, self.clause_active, True)
+
+    def fit(self, X, Y, shuffle=True):
+        if self.initialized == False:
+            self.initialize(X, Y)
+            self.initialized = True
+
+        if not np.array_equal(self.X_train, X):
+            self.encoded_X_train = self.clause_bank.prepare_X(X)
+            self.X_train = X.copy()
+
+        Ym = np.ascontiguousarray(Y).astype(np.uint32)
+
+        # Clauses are dropped based on their weights
+        self.clause_active = np.ones(self.number_of_clauses, dtype=np.uint32)
+        clause_score = np.zeros(self.number_of_clauses, dtype=np.int32)
+        for i in range(self.number_of_classes):
+            clause_score += np.abs(self.weight_banks[i].get_weights())
+        deactivate = np.random.choice(np.arange(self.number_of_clauses),
+                                      size=int(self.number_of_clauses * self.clause_drop_p),
+                                      p=clause_score / clause_score.sum())
+        for d in range(deactivate.shape[0]):
+            self.clause_active[deactivate[d]] = 0
+
+        # Literals are dropped based on their frequency
+        self.literal_active = (np.zeros(self.clause_bank.number_of_ta_chunks, dtype=np.uint32) | ~0).astype(np.uint32)
+        literal_clause_frequency = self.literal_clause_frequency()
+        if literal_clause_frequency.sum() > 0:
+            deactivate = np.random.choice(np.arange(self.clause_bank.number_of_literals),
+                                          size=int(self.clause_bank.number_of_literals * self.literal_drop_p),
+                                          p=literal_clause_frequency / literal_clause_frequency.sum())
+        else:
+            deactivate = np.random.choice(np.arange(self.clause_bank.number_of_literals),
+                                          size=int(self.clause_bank.number_of_literals * self.literal_drop_p))
+        for d in range(deactivate.shape[0]):
+            ta_chunk = deactivate[d] // 32
+            chunk_pos = deactivate[d] % 32
+            self.literal_active[ta_chunk] &= (~(1 << chunk_pos))
+
+        if not self.feature_negation:
+            for k in range(self.clause_bank.number_of_literals // 2, self.clause_bank.number_of_literals):
+                ta_chunk = k // 32
+                chunk_pos = k % 32
+                self.literal_active[ta_chunk] &= (~(1 << chunk_pos))
+
+        self.literal_active = self.literal_active.astype(np.uint32)
+
+        self.update_ps = np.empty(self.number_of_classes)
+
+        shuffled_index = np.arange(X.shape[0])
+        if shuffle:
+            np.random.shuffle(shuffled_index)
+
+        class_observed = np.zeros(self.number_of_classes, dtype=np.uint32)
+        example_indexes = np.zeros(self.number_of_classes, dtype=np.uint32)
+        example_counter = 0
+        for e in shuffled_index:
+            if self.output_balancing:
+                if class_observed[Ym[e]] == 0:
+                    example_indexes[Ym[e]] = e
+                    class_observed[Ym[e]] = 1
+                    example_counter += 1
+            else:
+                example_indexes[example_counter] = e
+                example_counter += 1
+
+            if example_counter == self.number_of_classes:
+                example_counter = 0
+
+                for i in range(self.number_of_classes):
+                    class_observed[i] = 0
+                    batch_example = example_indexes[i]
+                    self.update(Ym[batch_example], batch_example)
+        return
+
+    def predict(self, X):
+        if not np.array_equal(self.X_test, X):
+            self.encoded_X_test = self.clause_bank.prepare_X(X)
+            self.X_test = X.copy()
+
+        Y = np.ascontiguousarray(np.zeros(X.shape[0], dtype=np.uint32))
+
+        for e in range(X.shape[0]):
+            max_class_sum = -self.T
+            max_class = 0
+            clause_outputs = self.clause_bank.calculate_clause_outputs_predict(self.encoded_X_test, e)
+            for i in range(self.number_of_classes):
+                class_sum = np.dot(self.weight_banks[i].get_weights(), clause_outputs).astype(np.int32)
+                class_sum = np.clip(class_sum, -self.T, self.T)
+                if class_sum > max_class_sum:
+                    max_class_sum = class_sum
+                    max_class = i
+            Y[e] = max_class
+        return Y
+
+    def predict_individual(self, X):
+        if not np.array_equal(self.X_test, X):
+            self.encoded_X_test = self.clause_bank.prepare_X(X)
+            self.X_test = X.copy()
+
+        Y = np.ascontiguousarray(np.zeros((X.shape[0], self.number_of_classes), dtype=np.uint32))
+
+        for e in range(X.shape[0]):
+            clause_outputs = self.clause_bank.calculate_clause_outputs_predict(self.encoded_X_test, e)
+            for i in range(self.number_of_classes):
+                class_sum = np.dot(self.weight_banks[i].get_weights(), clause_outputs).astype(np.int32)
+                class_sum = np.clip(class_sum, -self.T, self.T)
+                Y[e, i] = (class_sum >= 0)
+        return Y
+
+    def clause_precision(self, the_class, positive_polarity, X, Y):
+        clause_outputs = self.transform(X)
+        weights = self.weight_banks[the_class].get_weights()
+        if positive_polarity == 0:
+            positive_clause_outputs = (weights >= 0)[:, np.newaxis].transpose() * clause_outputs
+            true_positive_clause_outputs = positive_clause_outputs[Y == the_class].sum(axis=0)
+            false_positive_clause_outputs = positive_clause_outputs[Y != the_class].sum(axis=0)
+        else:
+            positive_clause_outputs = (weights < 0)[:, np.newaxis].transpose() * clause_outputs
+            true_positive_clause_outputs = positive_clause_outputs[Y != the_class].sum(axis=0)
+            false_positive_clause_outputs = positive_clause_outputs[Y == the_class].sum(axis=0)
+
+        return np.where(true_positive_clause_outputs + false_positive_clause_outputs == 0, 0,
+                        1.0 * true_positive_clause_outputs / (
+                                    true_positive_clause_outputs + false_positive_clause_outputs))
+
+    def clause_recall(self, the_class, positive_polarity, X, Y):
+        clause_outputs = self.transform(X)
+        weights = self.weight_banks[the_class].get_weights()
+
+        if positive_polarity == 0:
+            positive_clause_outputs = (weights >= 0)[:, np.newaxis].transpose() * clause_outputs
+            true_positive_clause_outputs = positive_clause_outputs[Y == the_class].sum(axis=0) / \
+                                           Y[Y == the_class].shape[0]
+        else:
+            positive_clause_outputs = (weights < 0)[:, np.newaxis].transpose() * clause_outputs
+            true_positive_clause_outputs = positive_clause_outputs[Y != the_class].sum(axis=0) / \
+                                           Y[Y != the_class].shape[0]
+
+        return true_positive_clause_outputs
+
+    def get_weight(self, the_class, clause):
+        return self.weight_banks[the_class].get_weights()[clause]
+
+    def set_weight(self, the_class, clause, weight):
+        self.weight_banks[the_class].get_weights()[clause] = weight
+
+    def number_of_include_actions(self, clause):
+        return self.clause_bank.number_of_include_actions(clause)
diff --git a/tmu/models/classification/multichannel_classifier.py b/tmu/models/classification/multichannel_classifier.py
new file mode 100644
index 00000000..9b08f6f7
--- /dev/null
+++ b/tmu/models/classification/multichannel_classifier.py
@@ -0,0 +1,241 @@
+# Copyright (c) 2023 Ole-Christoffer Granmo
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+from tmu.clause_bank import ClauseBank
+from tmu.models.base import TMBasis
+from tmu.weight_bank import WeightBank
+import numpy as np
+import sys
+
+
+class TMMultiChannelClassifier(TMBasis):
+    def __init__(self, number_of_clauses, global_T, T, s, type_i_ii_ratio=1.0, platform='CPU', patch_dim=None,
+                 feature_negation=True, boost_true_positive_feedback=1, max_included_literals=None,
+                 number_of_state_bits_ta=8, weighted_clauses=False, clause_drop_p=0.0, literal_drop_p=0.0):
+        super().__init__(number_of_clauses, T, s, type_i_ii_ratio=type_i_ii_ratio, platform=platform,
+                         patch_dim=patch_dim, feature_negation=feature_negation,
+                         boost_true_positive_feedback=boost_true_positive_feedback,
+                         max_included_literals=max_included_literals, number_of_state_bits_ta=number_of_state_bits_ta,
+                         weighted_clauses=weighted_clauses, clause_drop_p=clause_drop_p, literal_drop_p=literal_drop_p)
+        self.global_T = global_T
+
+    def initialize(self, X, Y):
+        self.number_of_classes = int(np.max(Y) + 1)
+
+        if self.platform == 'CPU':
+            self.clause_bank = ClauseBank(X[0], self.number_of_clauses, self.number_of_state_bits_ta,
+                                          self.number_of_state_bits_ind, self.patch_dim)
+        elif self.platform == 'CUDA':
+            from tmu.clause_bank_cuda import ClauseBankCUDA
+            self.clause_bank = ClauseBankCUDA(X[0], self.number_of_clauses, self.number_of_state_bits_ta,
+                                              self.patch_dim)
+        else:
+            raise RuntimeError(f"Unknown platform of type: {self.platform}")
+
+        self.weight_banks = []
+        for i in range(self.number_of_classes):
+            self.weight_banks.append(
+                WeightBank(np.random.choice([-1, 1], size=self.number_of_clauses).astype(np.int32)))
+
+        self.X_train = {}
+        self.X_test = {}
+        for c in range(X.shape[0]):
+            self.X_train[c] = np.zeros(0, dtype=np.uint32)
+            self.X_test[c] = np.zeros(0, dtype=np.uint32)
+
+        self.encoded_X_train = {}
+        self.encoded_X_test = {}
+
+        if self.max_included_literals == None:
+            self.max_included_literals = self.clause_bank.number_of_literals
+
+    def fit(self, X, Y, shuffle=True):
+        if self.initialized == False:
+            self.initialize(X, Y)
+            self.initialized = True
+
+        for c in range(X.shape[0]):
+            if not np.array_equal(self.X_train[c], X[c]):
+                self.encoded_X_train[c] = self.clause_bank.prepare_X(X[c])
+                self.X_train[c] = X[c].copy()
+
+        Ym = np.ascontiguousarray(Y).astype(np.uint32)
+
+        # Clauses are dropped based on their weights
+        clause_active = np.ones(self.number_of_clauses, dtype=np.uint32)
+        clause_score = np.zeros(self.number_of_clauses, dtype=np.int32)
+        for i in range(self.number_of_classes):
+            clause_score += np.abs(self.weight_banks[i].get_weights())
+        deactivate = np.random.choice(np.arange(self.number_of_clauses),
+                                      size=int(self.number_of_clauses * self.clause_drop_p),
+                                      p=clause_score / clause_score.sum())
+        for d in range(deactivate.shape[0]):
+            clause_active[deactivate[d]] = 0
+
+        # Literals are dropped based on their frequency
+        literal_active = (np.zeros(self.clause_bank.number_of_ta_chunks, dtype=np.uint32) | ~0).astype(np.uint32)
+        literal_clause_frequency = self.literal_clause_frequency()
+        if literal_clause_frequency.sum() > 0:
+            deactivate = np.random.choice(np.arange(self.clause_bank.number_of_literals),
+                                          size=int(self.clause_bank.number_of_literals * self.literal_drop_p),
+                                          p=literal_clause_frequency / literal_clause_frequency.sum())
+        else:
+            deactivate = np.random.choice(np.arange(self.clause_bank.number_of_literals),
+                                          size=int(self.clause_bank.number_of_literals * self.literal_drop_p))
+        for d in range(deactivate.shape[0]):
+            ta_chunk = deactivate[d] // 32
+            chunk_pos = deactivate[d] % 32
+            literal_active[ta_chunk] &= (~(1 << chunk_pos))
+
+        if not self.feature_negation:
+            for k in range(self.clause_bank.number_of_literals // 2, self.clause_bank.number_of_literals):
+                ta_chunk = k // 32
+                chunk_pos = k % 32
+                literal_active[ta_chunk] &= (~(1 << chunk_pos))
+
+        literal_active = literal_active.astype(np.uint32)
+
+        local_class_sum = np.empty(X.shape[0], dtype=np.int32)
+
+        shuffled_index = np.arange(X.shape[1])
+        if shuffle:
+            np.random.shuffle(shuffled_index)
+
+        for e in shuffled_index:
+            target = Ym[e]
+
+            clause_outputs = []
+            for c in range(X.shape[0]):
+                clause_outputs.append(
+                    self.clause_bank.calculate_clause_outputs_update(literal_active, self.encoded_X_train[c], e).copy())
+
+            global_class_sum = 0
+            for c in range(X.shape[0]):
+                local_class_sum[c] = np.dot(clause_active * self.weight_banks[target].get_weights(),
+                                            clause_outputs[c]).astype(np.int32)
+                local_class_sum[c] = np.clip(local_class_sum[c], -self.T, self.T)
+                global_class_sum += local_class_sum[c]
+            global_class_sum = np.clip(global_class_sum, -self.global_T[target][0], self.global_T[target][1])
+            global_update_p = 1.0 * (self.global_T[target][1] - global_class_sum) / (
+                        self.global_T[target][0] + self.global_T[target][1])
+
+            for c in range(X.shape[0]):
+                local_update_p = 1.0 * (self.T - local_class_sum[c]) / (2 * self.T)
+                update_p = np.minimum(local_update_p, global_update_p)
+                self.clause_bank.type_i_feedback(update_p * self.type_i_p, self.s[target],
+                                                 self.boost_true_positive_feedback,
+                                                 clause_active * (self.weight_banks[target].get_weights() >= 0),
+                                                 literal_active, self.encoded_X_train[c], e)
+                self.clause_bank.type_ii_feedback(update_p * self.type_ii_p,
+                                                  clause_active * (self.weight_banks[target].get_weights() < 0),
+                                                  literal_active, self.encoded_X_train[c], e)
+                self.weight_banks[target].increment(clause_outputs[c], update_p, clause_active, True)
+
+            not_target = np.random.randint(self.number_of_classes)
+            while not_target == target:
+                not_target = np.random.randint(self.number_of_classes)
+
+            global_class_sum = 0.0
+            for c in range(X.shape[0]):
+                local_class_sum[c] = np.dot(clause_active * self.weight_banks[not_target].get_weights(),
+                                            clause_outputs[c]).astype(np.int32)
+                local_class_sum[c] = np.clip(local_class_sum[c], -self.T, self.T)
+                global_class_sum += local_class_sum[c]
+            global_class_sum = np.clip(global_class_sum, -self.global_T[not_target][0], self.global_T[not_target][1])
+            global_update_p = 1.0 * (self.global_T[not_target][0] + global_class_sum) / (
+                        self.global_T[not_target][0] + self.global_T[not_target][1])
+
+            for c in range(X.shape[0]):
+                local_update_p = 1.0 * (self.T + local_class_sum[c]) / (2 * self.T)
+                update_p = np.minimum(local_update_p, global_update_p)
+                self.clause_bank.type_i_feedback(update_p * self.type_i_p, self.s[not_target],
+                                                 self.boost_true_positive_feedback, self.max_included_literals,
+                                                 clause_active * (self.weight_banks[not_target].get_weights() < 0),
+                                                 literal_active, self.encoded_X_train[c], e)
+                self.clause_bank.type_ii_feedback(update_p * self.type_ii_p,
+                                                  clause_active * (self.weight_banks[not_target].get_weights() >= 0),
+                                                  literal_active, self.encoded_X_train[c], e)
+                self.weight_banks[not_target].decrement(clause_outputs[c], update_p, clause_active, True)
+        return
+
+    def predict(self, X):
+        for c in range(X.shape[0]):
+            if not np.array_equal(self.X_test[c], X[c]):
+                self.encoded_X_test[c] = self.clause_bank.prepare_X(X[c])
+                self.X_test[c] = X[c].copy()
+
+        Y = np.ascontiguousarray(np.zeros(X.shape[1], dtype=np.uint32))
+
+        for e in range(X.shape[1]):
+            max_class_sum = -sys.maxsize
+            max_class = 0
+
+            clause_outputs = []
+            for c in range(X.shape[0]):
+                clause_outputs.append(
+                    self.clause_bank.calculate_clause_outputs_predict(self.encoded_X_test[c], e).copy())
+
+            for i in range(self.number_of_classes):
+                global_class_sum = 1
+                for c in range(X.shape[0]):
+                    local_class_sum = np.dot(self.weight_banks[i].get_weights(), clause_outputs[c]).astype(np.int32)
+                    local_class_sum = np.clip(local_class_sum, -self.T, self.T)
+                    global_class_sum *= local_class_sum >= 0
+                global_class_sum = np.clip(global_class_sum, -self.global_T[i][0], self.global_T[i][1])
+
+                if global_class_sum > max_class_sum:
+                    max_class_sum = global_class_sum
+                    max_class = i
+            Y[e] = max_class
+
+        return Y
+
+    def clause_precision(self, the_class, positive_polarity, X, Y):
+        clause_outputs = self.transform(X)
+        weights = self.weight_banks[the_class].get_weights()
+        if positive_polarity == 0:
+            positive_clause_outputs = (weights >= 0)[:, np.newaxis].transpose() * clause_outputs
+            true_positive_clause_outputs = clause_outputs[Y == the_class].sum(axis=0)
+            false_positive_clause_outputs = clause_outputs[Y != the_class].sum(axis=0)
+        else:
+            positive_clause_outputs = (weights < 0)[:, np.newaxis].transpose() * clause_outputs
+            true_positive_clause_outputs = clause_outputs[Y != the_class].sum(axis=0)
+            false_positive_clause_outputs = clause_outputs[Y == the_class].sum(axis=0)
+
+        return np.where(true_positive_clause_outputs + false_positive_clause_outputs == 0, 0,
+                        1.0 * true_positive_clause_outputs / (
+                                    true_positive_clause_outputs + false_positive_clause_outputs))
+
+    def clause_recall(self, the_class, positive_polarity, X, Y):
+        clause_outputs = self.transform(X)
+        weights = self.weight_banks[the_class].get_weights()
+
+        if positive_polarity == 0:
+            positive_clause_outputs = (weights >= 0)[:, np.newaxis].transpose() * clause_outputs
+            true_positive_clause_outputs = positive_clause_outputs[Y == the_class].sum(axis=0)
+        else:
+            positive_clause_outputs = (weights < 0)[:, np.newaxis].transpose() * clause_outputs
+            true_positive_clause_outputs = positive_clause_outputs[Y != the_class].sum(axis=0)
+
+        return true_positive_clause_outputs / Y[Y == the_class].shape[0]
+
+    def get_weight(self, the_class, clause):
+        return self.weight_banks[the_class].get_weights()[clause]
+
+    def set_weight(self, the_class, clause, weight):
+        self.weight_banks[the_class].get_weights()[clause] = weight
diff --git a/tmu/models/classification/multitask_classifier.py b/tmu/models/classification/multitask_classifier.py
new file mode 100644
index 00000000..75d49e87
--- /dev/null
+++ b/tmu/models/classification/multitask_classifier.py
@@ -0,0 +1,245 @@
+# Copyright (c) 2023 Ole-Christoffer Granmo
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+from tmu.clause_bank import ClauseBank
+from tmu.models.base import TMBasis
+from tmu.weight_bank import WeightBank
+import numpy as np
+from scipy.sparse import csr_matrix
+
+
+class TMMultiTaskClassifier(TMBasis):
+    def __init__(self, number_of_clauses, T, s, confidence_driven_updating=False, type_i_ii_ratio=1.0,
+                 type_iii_feedback=False, focused_negative_sampling=False, output_balancing=False, d=200.0,
+                 platform='CPU', patch_dim=None, feature_negation=True, boost_true_positive_feedback=1,
+                 max_included_literals=None, number_of_state_bits_ta=8, number_of_state_bits_ind=8,
+                 weighted_clauses=False, clause_drop_p=0.0, literal_drop_p=0.0):
+        super().__init__(number_of_clauses, T, s, confidence_driven_updating=confidence_driven_updating,
+                         type_i_ii_ratio=type_i_ii_ratio, type_iii_feedback=type_iii_feedback,
+                         focused_negative_sampling=focused_negative_sampling, output_balancing=output_balancing, d=d,
+                         platform=platform, patch_dim=patch_dim, feature_negation=feature_negation,
+                         boost_true_positive_feedback=boost_true_positive_feedback,
+                         max_included_literals=max_included_literals, number_of_state_bits_ta=number_of_state_bits_ta,
+                         number_of_state_bits_ind=number_of_state_bits_ind, weighted_clauses=weighted_clauses,
+                         clause_drop_p=clause_drop_p, literal_drop_p=literal_drop_p)
+
+    def initialize(self, X, Y):
+        self.number_of_classes = len(X)
+
+        if self.platform == 'CPU':
+            self.clause_bank = ClauseBank(X[0], self.number_of_clauses, self.number_of_state_bits_ta,
+                                          self.number_of_state_bits_ind, self.patch_dim)
+        elif self.platform == 'CUDA':
+            from tmu.clause_bank_cuda import ClauseBankCUDA
+            self.clause_bank = ClauseBankCUDA(X[0], self.number_of_clauses, self.number_of_state_bits_ta,
+                                              self.patch_dim)
+        else:
+            raise RuntimeError(f"Unknown platform of type: {self.platform}")
+
+        self.weight_banks = []
+        for i in range(self.number_of_classes):
+            self.weight_banks.append(
+                WeightBank(np.random.choice([-1, 1], size=self.number_of_clauses).astype(np.int32)))
+
+        if self.max_included_literals == None:
+            self.max_included_literals = self.clause_bank.number_of_literals
+
+    def fit(self, X, Y, shuffle=True):
+        if self.initialized == False:
+            self.initialize(X, Y)
+            self.initialized = True
+
+        X_csr = {}
+        for i in range(self.number_of_classes):
+            X_csr[i] = csr_matrix(X[i].reshape(X[i].shape[0], -1))
+
+        # Clauses are dropped based on their weights
+        self.clause_active = np.ones(self.number_of_clauses, dtype=np.uint32)
+        clause_score = np.zeros(self.number_of_clauses, dtype=np.int32)
+        for i in range(self.number_of_classes):
+            clause_score += np.abs(self.weight_banks[i].get_weights())
+        deactivate = np.random.choice(np.arange(self.number_of_clauses),
+                                      size=int(self.number_of_clauses * self.clause_drop_p),
+                                      p=clause_score / clause_score.sum())
+        for d in range(deactivate.shape[0]):
+            self.clause_active[deactivate[d]] = 0
+
+        # Literals are dropped based on their frequency
+        self.literal_active = (np.zeros(self.clause_bank.number_of_ta_chunks, dtype=np.uint32) | ~0).astype(np.uint32)
+        literal_clause_frequency = self.literal_clause_frequency()
+        if literal_clause_frequency.sum() > 0:
+            deactivate = np.random.choice(np.arange(self.clause_bank.number_of_literals),
+                                          size=int(self.clause_bank.number_of_literals * self.literal_drop_p),
+                                          p=literal_clause_frequency / literal_clause_frequency.sum())
+        else:
+            deactivate = np.random.choice(np.arange(self.clause_bank.number_of_literals),
+                                          size=int(self.clause_bank.number_of_literals * self.literal_drop_p))
+        for d in range(deactivate.shape[0]):
+            ta_chunk = deactivate[d] // 32
+            chunk_pos = deactivate[d] % 32
+            self.literal_active[ta_chunk] &= (~(1 << chunk_pos))
+
+        if not self.feature_negation:
+            for k in range(self.clause_bank.number_of_literals // 2, self.clause_bank.number_of_literals):
+                ta_chunk = k // 32
+                chunk_pos = k % 32
+                self.literal_active[ta_chunk] &= (~(1 << chunk_pos))
+
+        self.literal_active = self.literal_active.astype(np.uint32)
+
+        shuffled_index = np.arange(X[0].shape[0])
+        if shuffle:
+            np.random.shuffle(shuffled_index)
+
+        class_index = np.arange(self.number_of_classes, dtype=np.uint32)
+        for e in shuffled_index:
+            np.random.shuffle(class_index)
+
+            average_absolute_weights = np.zeros(self.number_of_clauses, dtype=np.float32)
+            for i in class_index:
+                average_absolute_weights += np.absolute(self.weight_banks[i].get_weights())
+            average_absolute_weights /= self.number_of_classes
+            update_clause = np.random.random(self.number_of_clauses) <= (
+                        self.T - np.clip(average_absolute_weights, 0, self.T)) / self.T
+
+            for i in class_index:
+                encoded_X = self.clause_bank.prepare_X(X_csr[i][e, :].toarray())
+                clause_outputs = self.clause_bank.calculate_clause_outputs_update(self.literal_active, encoded_X, 0)
+
+                class_sum = np.dot(self.clause_active * self.weight_banks[i].get_weights(), clause_outputs).astype(
+                    np.int32)
+                class_sum = np.clip(class_sum, -self.T, self.T)
+
+                type_iii_feedback_selection = np.random.choice(2)
+
+                if Y[i][e] == 1:
+                    if self.confidence_driven_updating:
+                        update_p = 1.0 * (self.T - np.absolute(class_sum)) / self.T
+                    else:
+                        update_p = (self.T - class_sum) / (2 * self.T)
+
+                    self.clause_bank.type_i_feedback(update_p * self.type_i_p, self.s,
+                                                     self.boost_true_positive_feedback, self.max_included_literals,
+                                                     update_clause * self.clause_active * (
+                                                                 self.weight_banks[i].get_weights() >= 0),
+                                                     self.literal_active, encoded_X, 0)
+                    self.clause_bank.type_ii_feedback(update_p * self.type_ii_p, update_clause * self.clause_active * (
+                                self.weight_banks[i].get_weights() < 0), self.literal_active, encoded_X, 0)
+                    self.weight_banks[i].increment(clause_outputs, update_p, update_clause * self.clause_active, True)
+                    if self.type_iii_feedback and type_iii_feedback_selection == 0:
+                        self.clause_bank.type_iii_feedback(update_p, self.d, update_clause * self.clause_active * (
+                                    self.weight_banks[i].get_weights() >= 0), self.literal_active, encoded_X, 0, 1)
+                        self.clause_bank.type_iii_feedback(update_p, self.d, update_clause * self.clause_active * (
+                                    self.weight_banks[i].get_weights() < 0), self.literal_active, encoded_X, 0, 0)
+                else:
+                    if self.confidence_driven_updating:
+                        update_p = 1.0 * (self.T - np.absolute(class_sum)) / self.T
+                    else:
+                        update_p = (self.T + class_sum) / (2 * self.T)
+
+                    self.clause_bank.type_i_feedback(update_p * self.type_i_p, self.s,
+                                                     self.boost_true_positive_feedback, self.max_included_literals,
+                                                     update_clause * self.clause_active * (
+                                                                 self.weight_banks[i].get_weights() < 0),
+                                                     self.literal_active, encoded_X, 0)
+                    self.clause_bank.type_ii_feedback(update_p * self.type_ii_p,
+                                                      update_clause * update_clause * self.clause_active * (
+                                                                  self.weight_banks[i].get_weights() >= 0),
+                                                      self.literal_active, encoded_X, 0)
+                    self.weight_banks[i].decrement(clause_outputs, update_p, self.clause_active, True)
+                    if self.type_iii_feedback and type_iii_feedback_selection == 1:
+                        self.clause_bank.type_iii_feedback(update_p, self.d, update_clause * self.clause_active * (
+                                    self.weight_banks[i].get_weights() < 0), self.literal_active, encoded_X, 0, 1)
+                        self.clause_bank.type_iii_feedback(update_p, self.d, update_clause * self.clause_active * (
+                                    self.weight_banks[i].get_weights() >= 0), self.literal_active, encoded_X, 0, 0)
+        return
+
+    def predict(self, X):
+        Y = {}
+        for i in range(len(X)):
+            Y[i] = np.ascontiguousarray(np.zeros((X[0].shape[0]), dtype=np.uint32))
+
+        X_csr = {}
+        for i in range(self.number_of_classes):
+            X_csr[i] = csr_matrix(X[i].reshape(X[i].shape[0], -1))
+
+        for i in range(self.number_of_classes):
+            for e in range(X[i].shape[0]):
+                encoded_X = self.clause_bank.prepare_X(X_csr[i][e, :].toarray())
+                clause_outputs = self.clause_bank.calculate_clause_outputs_predict(encoded_X, 0)
+                class_sum = np.dot(self.weight_banks[i].get_weights(), clause_outputs).astype(np.int32)
+                Y[i][e] = (class_sum >= 0)
+        return Y
+
+    def predict_exclusive(self, X):
+        if not np.array_equal(self.X_test, X):
+            self.encoded_X_test = self.clause_bank.prepare_X(X)
+            self.X_test = X.copy()
+
+        Y = np.ascontiguousarray(np.zeros(X.shape[0], dtype=np.uint32))
+
+        for e in range(X.shape[0]):
+            max_class_sum = -self.T
+            max_class = 0
+            clause_outputs = self.clause_bank.calculate_clause_outputs_predict(self.encoded_X_test, e)
+            for i in range(self.number_of_classes):
+                class_sum = np.dot(self.weight_banks[i].get_weights(), clause_outputs).astype(np.int32)
+                class_sum = np.clip(class_sum, -self.T, self.T)
+                if class_sum > max_class_sum:
+                    max_class_sum = class_sum
+                    max_class = i
+            Y[e] = max_class
+        return Y
+
+    def clause_precision(self, the_class, positive_polarity, X, Y):
+        clause_outputs = self.transform(X)
+        weights = self.weight_banks[the_class].get_weights()
+        if positive_polarity == 0:
+            positive_clause_outputs = (weights >= 0)[:, np.newaxis].transpose() * clause_outputs
+            true_positive_clause_outputs = clause_outputs[Y == the_class].sum(axis=0)
+            false_positive_clause_outputs = clause_outputs[Y != the_class].sum(axis=0)
+        else:
+            positive_clause_outputs = (weights < 0)[:, np.newaxis].transpose() * clause_outputs
+            true_positive_clause_outputs = clause_outputs[Y != the_class].sum(axis=0)
+            false_positive_clause_outputs = clause_outputs[Y == the_class].sum(axis=0)
+
+        return np.where(true_positive_clause_outputs + false_positive_clause_outputs == 0, 0,
+                        1.0 * true_positive_clause_outputs / (
+                                    true_positive_clause_outputs + false_positive_clause_outputs))
+
+    def clause_recall(self, the_class, positive_polarity, X, Y):
+        clause_outputs = self.transform(X)
+        weights = self.weight_banks[the_class].get_weights()
+
+        if positive_polarity == 0:
+            positive_clause_outputs = (weights >= 0)[:, np.newaxis].transpose() * clause_outputs
+            true_positive_clause_outputs = positive_clause_outputs[Y == the_class].sum(axis=0)
+        else:
+            positive_clause_outputs = (weights < 0)[:, np.newaxis].transpose() * clause_outputs
+            true_positive_clause_outputs = positive_clause_outputs[Y != the_class].sum(axis=0)
+
+        return true_positive_clause_outputs / Y[Y == the_class].shape[0]
+
+    def get_weight(self, the_class, clause):
+        return self.weight_banks[the_class].get_weights()[clause]
+
+    def get_weights(self, the_class):
+        return self.weight_banks[the_class].get_weights()
+
+    def set_weight(self, the_class, clause, weight):
+        self.weight_banks[the_class].get_weights()[clause] = weight
diff --git a/tmu/models/classification/one_vs_one_classifier.py b/tmu/models/classification/one_vs_one_classifier.py
new file mode 100644
index 00000000..1a41c543
--- /dev/null
+++ b/tmu/models/classification/one_vs_one_classifier.py
@@ -0,0 +1,192 @@
+# Copyright (c) 2023 Ole-Christoffer Granmo
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+from tmu.clause_bank import ClauseBank
+from tmu.models.base import TMBasis
+from tmu.weight_bank import WeightBank
+import numpy as np
+
+
+class TMOneVsOneClassifier(TMBasis):
+    def __init__(self, number_of_clauses, T, s, type_i_ii_ratio=1.0, platform='CPU', patch_dim=None,
+                 feature_negation=True, boost_true_positive_feedback=1, max_included_literals=None,
+                 number_of_state_bits_ta=8, weighted_clauses=False, clause_drop_p=0.0, literal_drop_p=0.0):
+        super().__init__(number_of_clauses, T, s, type_i_ii_ratio=type_i_ii_ratio, platform=platform,
+                         patch_dim=patch_dim, feature_negation=feature_negation,
+                         boost_true_positive_feedback=boost_true_positive_feedback,
+                         max_included_literals=max_included_literals, number_of_state_bits_ta=number_of_state_bits_ta,
+                         weighted_clauses=weighted_clauses, clause_drop_p=clause_drop_p, literal_drop_p=literal_drop_p)
+
+    def initialize(self, X, Y):
+        self.number_of_classes = int(np.max(Y) + 1)
+        self.number_of_outputs = self.number_of_classes * (self.number_of_classes - 1)
+
+        if self.platform == 'CPU':
+            self.clause_bank = ClauseBank(X, self.number_of_clauses, self.number_of_state_bits_ta,
+                                          self.number_of_state_bits_ind, self.patch_dim)
+        elif self.platform == 'CUDA':
+            from tmu.clause_bank_cuda import ClauseBankCUDA
+            self.clause_bank = ClauseBankCUDA(X, self.number_of_clauses, self.number_of_state_bits_ta, self.patch_dim)
+        else:
+            raise RuntimeError(f"Unknown platform of type: {self.platform}")
+
+        self.weight_banks = []
+        for i in range(self.number_of_outputs):
+            self.weight_banks.append(WeightBank(np.ones(self.number_of_clauses).astype(np.int32)))
+
+        if self.max_included_literals == None:
+            self.max_included_literals = self.clause_bank.number_of_literals
+
+    def fit(self, X, Y, shuffle=True):
+        if self.initialized == False:
+            self.initialize(X, Y)
+            self.initialized = True
+
+        if not np.array_equal(self.X_train, X):
+            self.encoded_X_train = self.clause_bank.prepare_X(X)
+            self.X_train = X.copy()
+
+        Ym = np.ascontiguousarray(Y).astype(np.uint32)
+
+        clause_active = np.ascontiguousarray(
+            np.random.choice(2, self.number_of_clauses, p=[self.clause_drop_p, 1.0 - self.clause_drop_p]).astype(
+                np.int32))
+        literal_active = (np.zeros(self.clause_bank.number_of_ta_chunks, dtype=np.uint32) | ~0).astype(np.uint32)
+        if not self.feature_negation:
+            for k in range(self.clause_bank.number_of_literals // 2, self.clause_bank.number_of_literals):
+                ta_chunk = k // 32
+                chunk_pos = k % 32
+                literal_active[ta_chunk] &= (~(1 << chunk_pos))
+
+        literal_active = literal_active.astype(np.uint32)
+
+        shuffled_index = np.arange(X.shape[0])
+        if shuffle:
+            np.random.shuffle(shuffled_index)
+
+        for e in shuffled_index:
+            clause_outputs = self.clause_bank.calculate_clause_outputs_update(literal_active, self.encoded_X_train, e)
+
+            target = Ym[e]
+            not_target = np.random.randint(self.number_of_classes)
+            while not_target == target:
+                not_target = np.random.randint(self.number_of_classes)
+
+            output = target * (self.number_of_classes - 1) + not_target - (not_target > target)
+
+            class_sum = np.dot(clause_active * self.weight_banks[output].get_weights(), clause_outputs).astype(np.int32)
+            class_sum = np.clip(class_sum, -self.T, self.T)
+            update_p = (self.T - class_sum) / (2 * self.T)
+
+            self.clause_bank.type_i_feedback(update_p * self.type_i_p, self.s, self.boost_true_positive_feedback,
+                                             self.max_included_literals,
+                                             clause_active * (self.weight_banks[output].get_weights() >= 0),
+                                             literal_active, self.encoded_X_train, e)
+            self.clause_bank.type_ii_feedback(update_p * self.type_ii_p,
+                                              clause_active * (self.weight_banks[output].get_weights() < 0),
+                                              literal_active, self.encoded_X_train, e)
+            self.weight_banks[output].increment(clause_outputs, update_p, clause_active, True)
+
+            output = not_target * (self.number_of_classes - 1) + target - (target > not_target)
+
+            class_sum = np.dot(clause_active * self.weight_banks[output].get_weights(), clause_outputs).astype(np.int32)
+            class_sum = np.clip(class_sum, -self.T, self.T)
+            update_p = (self.T + class_sum) / (2 * self.T)
+
+            self.clause_bank.type_i_feedback(update_p * self.type_i_p, self.s, self.boost_true_positive_feedback,
+                                             self.max_included_literals,
+                                             clause_active * (self.weight_banks[output].get_weights() < 0),
+                                             literal_active, self.encoded_X_train, e)
+            self.clause_bank.type_ii_feedback(update_p * self.type_ii_p,
+                                              clause_active * (self.weight_banks[output].get_weights() >= 0),
+                                              literal_active, self.encoded_X_train, e)
+            self.weight_banks[output].decrement(clause_outputs, update_p, clause_active, True)
+        return
+
+    def predict(self, X):
+        if not np.array_equal(self.X_test, X):
+            self.encoded_X_test = self.clause_bank.prepare_X(X)
+            self.X_test = X.copy()
+
+        Y = np.ascontiguousarray(np.zeros(X.shape[0], dtype=np.uint32))
+
+        for e in range(X.shape[0]):
+            clause_outputs = self.clause_bank.calculate_clause_outputs_predict(self.encoded_X_test, e)
+
+            max_class_sum = -self.T * self.number_of_classes
+            max_class = 0
+            for i in range(self.number_of_classes):
+                class_sum = 0
+                for output in range(i * (self.number_of_classes - 1), (i + 1) * (self.number_of_classes - 1)):
+                    output_sum = np.dot(self.weight_banks[output].get_weights(), clause_outputs).astype(np.int32)
+                    output_sum = np.clip(output_sum, -self.T, self.T)
+                    class_sum += output_sum
+
+                if class_sum > max_class_sum:
+                    max_class_sum = class_sum
+                    max_class = i
+            Y[e] = max_class
+        return Y
+
+    def clause_precision(self, the_class, positive_polarity, X, Y):
+        clause_outputs = self.transform(X)
+        precision = np.zeros((self.number_of_classes - 1, self.number_of_clauses))
+        for i in range(self.number_of_classes - 1):
+            other_class = i + (i >= the_class)
+            output = the_class * (self.number_of_classes - 1) + i
+            weights = self.weight_banks[output].get_weights()
+            if positive_polarity:
+                positive_clause_outputs = (weights >= 0)[:, np.newaxis].transpose() * clause_outputs
+                true_positive_clause_outputs = positive_clause_outputs[Y == the_class].sum(axis=0)
+                false_positive_clause_outputs = positive_clause_outputs[Y == other_class].sum(axis=0)
+                precision[i] = np.where(true_positive_clause_outputs + false_positive_clause_outputs == 0, 0,
+                                        true_positive_clause_outputs / (
+                                                    true_positive_clause_outputs + false_positive_clause_outputs))
+            else:
+                positive_clause_outputs = (weights < 0)[:, np.newaxis].transpose() * clause_outputs
+                true_positive_clause_outputs = positive_clause_outputs[Y == other_class].sum(axis=0)
+                false_positive_clause_outputs = positive_clause_outputs[Y == the_class].sum(axis=0)
+                precision[i] = np.where(true_positive_clause_outputs + false_positive_clause_outputs == 0, 0,
+                                        true_positive_clause_outputs / (
+                                                    true_positive_clause_outputs + false_positive_clause_outputs))
+
+        return precision
+
+    def clause_recall(self, the_class, positive_polarity, X, Y):
+        clause_outputs = self.transform(X)
+        recall = np.zeros((self.number_of_classes - 1, self.number_of_clauses))
+        for i in range(self.number_of_classes - 1):
+            other_class = i + (i >= the_class)
+            output = the_class * (self.number_of_classes - 1) + i
+            weights = self.weight_banks[output].get_weights()
+            if positive_polarity:
+                positive_clause_outputs = (weights >= 0)[:, np.newaxis].transpose() * clause_outputs
+                true_positive_clause_outputs = positive_clause_outputs[Y == the_class].sum(axis=0)
+                recall[i] = true_positive_clause_outputs / Y[Y == the_class].shape[0]
+            else:
+                positive_clause_outputs = (weights < 0)[:, np.newaxis].transpose() * clause_outputs
+                true_positive_clause_outputs = positive_clause_outputs[Y == other_class].sum(axis=0)
+                recall[i] = true_positive_clause_outputs / Y[Y == other_class].shape[0]
+        return recall
+
+    def get_weight(self, output, clause):
+        return self.weight_banks[output].get_weights()[clause]
+
+    def set_weight(self, output, weight):
+        self.weight_banks[output].get_weights()[output] = weight
diff --git a/tmu/models/classification/vanilla_classifier.py b/tmu/models/classification/vanilla_classifier.py
new file mode 100644
index 00000000..e0906de8
--- /dev/null
+++ b/tmu/models/classification/vanilla_classifier.py
@@ -0,0 +1,307 @@
+# Copyright (c) 2023 Ole-Christoffer Granmo
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+from tmu.clause_bank import ClauseBank
+from tmu.models.base import TMBasis
+from tmu.weight_bank import WeightBank
+import numpy as np
+import tmu.tools
+
+
+class TMClassifier(TMBasis):
+    def __init__(self, number_of_clauses, T, s, confidence_driven_updating=False, type_i_ii_ratio=1.0,
+                 type_iii_feedback=False, d=200.0, platform='CPU', patch_dim=None, feature_negation=True,
+                 boost_true_positive_feedback=1, max_included_literals=None, number_of_state_bits_ta=8,
+                 number_of_state_bits_ind=8, weighted_clauses=False, clause_drop_p=0.0, literal_drop_p=0.0,
+                 batch_size=100, incremental=True):
+        super().__init__(number_of_clauses, T, s, confidence_driven_updating=confidence_driven_updating,
+                         type_i_ii_ratio=type_i_ii_ratio, type_iii_feedback=type_iii_feedback, d=d, platform=platform,
+                         patch_dim=patch_dim, feature_negation=feature_negation,
+                         boost_true_positive_feedback=boost_true_positive_feedback,
+                         max_included_literals=max_included_literals, number_of_state_bits_ta=number_of_state_bits_ta,
+                         number_of_state_bits_ind=number_of_state_bits_ind, weighted_clauses=weighted_clauses,
+                         clause_drop_p=clause_drop_p, literal_drop_p=literal_drop_p, batch_size=batch_size,
+                         incremental=incremental)
+
+    def initialize(self, X, Y):
+        self.number_of_classes = int(np.max(Y) + 1)
+
+        self.weight_banks = []
+        for i in range(self.number_of_classes):
+            self.weight_banks.append(WeightBank(np.concatenate((np.ones(self.number_of_clauses // 2, dtype=np.int32),
+                                                                -1 * np.ones(self.number_of_clauses // 2,
+                                                                             dtype=np.int32)))))
+
+        self.clause_banks = []
+        if self.platform == 'CPU':
+            for i in range(self.number_of_classes):
+                self.clause_banks.append(
+                    ClauseBank(X, self.number_of_clauses, self.number_of_state_bits_ta, self.number_of_state_bits_ind,
+                               self.patch_dim, batch_size=self.batch_size, incremental=self.incremental))
+        elif self.platform == 'CUDA':
+            from tmu.clause_bank_cuda import ClauseBankCUDA
+            for i in range(self.number_of_classes):
+                self.clause_banks.append(
+                    ClauseBankCUDA(X, self.number_of_clauses, self.number_of_state_bits_ta, self.patch_dim))
+        else:
+            raise RuntimeError(f"Unknown platform of type: {self.platform}")
+
+        if self.max_included_literals == None:
+            self.max_included_literals = self.clause_banks[0].number_of_literals
+
+        self.positive_clauses = np.concatenate((np.ones(self.number_of_clauses // 2, dtype=np.int32),
+                                                np.zeros(self.number_of_clauses // 2, dtype=np.int32)))
+        self.negative_clauses = np.concatenate((np.zeros(self.number_of_clauses // 2, dtype=np.int32),
+                                                np.ones(self.number_of_clauses // 2, dtype=np.int32)))
+
+    def fit(self, X, Y, shuffle=True):
+        if self.initialized == False:
+            self.initialize(X, Y)
+            self.initialized = True
+
+        if not np.array_equal(self.X_train, X):
+            self.encoded_X_train = self.clause_banks[0].prepare_X(X)
+            self.X_train = X.copy()
+
+        Ym = np.ascontiguousarray(Y).astype(np.uint32)
+
+        clause_active = []
+        for i in range(self.number_of_classes):
+            # Clauses are dropped based on their weights
+            class_clause_active = np.ascontiguousarray(np.ones(self.number_of_clauses, dtype=np.int32))
+            clause_score = np.abs(self.weight_banks[i].get_weights())
+            deactivate = np.random.choice(np.arange(self.number_of_clauses),
+                                          size=int(self.number_of_clauses * self.clause_drop_p),
+                                          p=clause_score / clause_score.sum())
+            for d in range(deactivate.shape[0]):
+                class_clause_active[deactivate[d]] = 0
+            clause_active.append(class_clause_active)
+
+        # Literals are dropped based on their frequency
+        literal_active = (np.zeros(self.clause_banks[0].number_of_ta_chunks, dtype=np.uint32) | ~0).astype(np.uint32)
+        literal_clause_frequency = self.literal_clause_frequency()
+        if literal_clause_frequency.sum() > 0:
+            deactivate = np.random.choice(np.arange(self.clause_banks[0].number_of_literals),
+                                          size=int(self.clause_banks[0].number_of_literals * self.literal_drop_p),
+                                          p=literal_clause_frequency / literal_clause_frequency.sum())
+        else:
+            deactivate = np.random.choice(np.arange(self.clause_banks[0].number_of_literals),
+                                          size=int(self.clause_banks[0].number_of_literals * self.literal_drop_p))
+        for d in range(deactivate.shape[0]):
+            ta_chunk = deactivate[d] // 32
+            chunk_pos = deactivate[d] % 32
+            literal_active[ta_chunk] &= (~(1 << chunk_pos))
+
+        if not self.feature_negation:
+            for k in range(self.clause_banks[0].number_of_literals // 2, self.clause_banks[0].number_of_literals):
+                ta_chunk = k // 32
+                chunk_pos = k % 32
+                literal_active[ta_chunk] &= (~(1 << chunk_pos))
+        literal_active = literal_active.astype(np.uint32)
+
+        shuffled_index = np.arange(X.shape[0])
+        if shuffle:
+            np.random.shuffle(shuffled_index)
+
+        for e in shuffled_index:
+            target = Ym[e]
+
+            clause_outputs = self.clause_banks[target].calculate_clause_outputs_update(literal_active,
+                                                                                       self.encoded_X_train, e)
+            class_sum = np.dot(clause_active[target] * self.weight_banks[target].get_weights(), clause_outputs).astype(
+                np.int32)
+            class_sum = np.clip(class_sum, -self.T, self.T)
+
+            if self.confidence_driven_updating:
+                update_p = 1.0 * (self.T - np.absolute(class_sum)) / self.T
+            else:
+                update_p = (self.T - class_sum) / (2 * self.T)
+
+            if self.weighted_clauses:
+                self.weight_banks[target].increment(clause_outputs, update_p, clause_active[target], False)
+            self.clause_banks[target].type_i_feedback(update_p * self.type_i_p, self.s,
+                                                      self.boost_true_positive_feedback, self.max_included_literals,
+                                                      clause_active[target] * self.positive_clauses, literal_active,
+                                                      self.encoded_X_train, e)
+            self.clause_banks[target].type_ii_feedback(update_p * self.type_ii_p,
+                                                       clause_active[target] * self.negative_clauses, literal_active,
+                                                       self.encoded_X_train, e)
+            if self.type_iii_feedback:
+                self.clause_banks[target].type_iii_feedback(update_p, self.d,
+                                                            clause_active[target] * self.positive_clauses,
+                                                            literal_active, self.encoded_X_train, e, 1)
+                self.clause_banks[target].type_iii_feedback(update_p, self.d,
+                                                            clause_active[target] * self.negative_clauses,
+                                                            literal_active, self.encoded_X_train, e, 0)
+
+            not_target = np.random.randint(self.number_of_classes)
+            while not_target == target:
+                not_target = np.random.randint(self.number_of_classes)
+
+            clause_outputs = self.clause_banks[not_target].calculate_clause_outputs_update(literal_active,
+                                                                                           self.encoded_X_train, e)
+            class_sum = np.dot(clause_active[not_target] * self.weight_banks[not_target].get_weights(),
+                               clause_outputs).astype(np.int32)
+            class_sum = np.clip(class_sum, -self.T, self.T)
+
+            if self.confidence_driven_updating:
+                update_p = 1.0 * (self.T - np.absolute(class_sum)) / self.T
+            else:
+                update_p = (self.T + class_sum) / (2 * self.T)
+
+            if self.weighted_clauses:
+                self.weight_banks[not_target].decrement(clause_outputs, update_p, clause_active[not_target], False)
+            self.clause_banks[not_target].type_i_feedback(update_p * self.type_i_p, self.s,
+                                                          self.boost_true_positive_feedback, self.max_included_literals,
+                                                          clause_active[not_target] * self.negative_clauses,
+                                                          literal_active, self.encoded_X_train, e)
+            self.clause_banks[not_target].type_ii_feedback(update_p * self.type_ii_p,
+                                                           clause_active[not_target] * self.positive_clauses,
+                                                           literal_active, self.encoded_X_train, e)
+            if self.type_iii_feedback:
+                self.clause_banks[not_target].type_iii_feedback(update_p, self.d,
+                                                                clause_active[not_target] * self.negative_clauses,
+                                                                literal_active, self.encoded_X_train, e, 1)
+                self.clause_banks[not_target].type_iii_feedback(update_p, self.d,
+                                                                clause_active[not_target] * self.positive_clauses,
+                                                                literal_active, self.encoded_X_train, e, 0)
+        return
+
+    def predict(self, X):
+        if not np.array_equal(self.X_test, X):
+            self.encoded_X_test = self.clause_banks[0].prepare_X(X)
+            self.X_test = X.copy()
+
+        Y = np.ascontiguousarray(np.zeros(X.shape[0], dtype=np.uint32))
+        for e in range(X.shape[0]):
+            max_class_sum = -self.T
+            max_class = 0
+            for i in range(self.number_of_classes):
+                class_sum = np.dot(self.weight_banks[i].get_weights(),
+                                   self.clause_banks[i].calculate_clause_outputs_predict(self.encoded_X_test,
+                                                                                         e)).astype(np.int32)
+                class_sum = np.clip(class_sum, -self.T, self.T)
+                if class_sum > max_class_sum:
+                    max_class_sum = class_sum
+                    max_class = i
+            Y[e] = max_class
+        return Y
+
+    def transform(self, X):
+        encoded_X = self.clause_banks[0].prepare_X(X)
+        transformed_X = np.empty((X.shape[0], self.number_of_classes, self.number_of_clauses), dtype=np.uint32)
+        for e in range(X.shape[0]):
+            for i in range(self.number_of_classes):
+                transformed_X[e, i, :] = self.clause_banks[i].calculate_clause_outputs_predict(encoded_X, e)
+        return transformed_X.reshape((X.shape[0], self.number_of_classes * self.number_of_clauses))
+
+    def transform_patchwise(self, X):
+        encoded_X = tmu.tools.encode(X, X.shape[0], self.number_of_patches, self.number_of_ta_chunks, self.dim,
+                                     self.patch_dim, 0)
+        transformed_X = np.empty(
+            (X.shape[0], self.number_of_classes, self.number_of_clauses // 2 * self.number_of_patches), dtype=np.uint32)
+        for e in range(X.shape[0]):
+            for i in range(self.number_of_classes):
+                transformed_X[e, i, :] = self.clause_bank[i].calculate_clause_outputs_patchwise(encoded_X, e)
+        return transformed_X.reshape(
+            (X.shape[0], self.number_of_classes * self.number_of_clauses, self.number_of_patches))
+
+    def literal_clause_frequency(self):
+        clause_active = np.ones(self.number_of_clauses, dtype=np.uint32)
+        literal_frequency = np.zeros(self.clause_banks[0].number_of_literals, dtype=np.uint32)
+        for i in range(self.number_of_classes):
+            literal_frequency += self.clause_banks[i].calculate_literal_clause_frequency(clause_active)
+        return literal_frequency
+
+    def literal_importance(self, the_class, negated_features=False, negative_polarity=False):
+        literal_frequency = np.zeros(self.clause_banks[0].number_of_literals, dtype=np.uint32)
+        if negated_features:
+            if negative_polarity:
+                literal_frequency[self.clause_banks[the_class].number_of_literals // 2:] += self.clause_banks[
+                                                                                                the_class].calculate_literal_clause_frequency(
+                    self.negative_clauses)[self.clause_banks[the_class].number_of_literals // 2:]
+            else:
+                literal_frequency[self.clause_banks[the_class].number_of_literals // 2:] += self.clause_banks[
+                                                                                                the_class].calculate_literal_clause_frequency(
+                    self.positive_clauses)[self.clause_banks[the_class].number_of_literals // 2:]
+        else:
+            if negative_polarity:
+                literal_frequency[:self.clause_banks[the_class].number_of_literals // 2] += self.clause_banks[
+                                                                                                the_class].calculate_literal_clause_frequency(
+                    self.negative_clauses)[:self.clause_banks[the_class].number_of_literals // 2]
+            else:
+                literal_frequency[:self.clause_banks[the_class].number_of_literals // 2] += self.clause_banks[
+                                                                                                the_class].calculate_literal_clause_frequency(
+                    self.positive_clauses)[:self.clause_banks[the_class].number_of_literals // 2]
+
+        return literal_frequency
+
+    def clause_precision(self, the_class, polarity, X, Y):
+        clause_outputs = self.transform(X).reshape(X.shape[0], self.number_of_classes, 2, self.number_of_clauses // 2)[
+                         :, the_class, polarity, :]
+        if polarity == 0:
+            true_positive_clause_outputs = clause_outputs[Y == the_class].sum(axis=0)
+            false_positive_clause_outputs = clause_outputs[Y != the_class].sum(axis=0)
+        else:
+            true_positive_clause_outputs = clause_outputs[Y != the_class].sum(axis=0)
+            false_positive_clause_outputs = clause_outputs[Y == the_class].sum(axis=0)
+        return np.where(true_positive_clause_outputs + false_positive_clause_outputs == 0, 0,
+                        true_positive_clause_outputs / (true_positive_clause_outputs + false_positive_clause_outputs))
+
+    def clause_recall(self, the_class, polarity, X, Y):
+        clause_outputs = self.transform(X).reshape(X.shape[0], self.number_of_classes, 2, self.number_of_clauses // 2)[
+                         :, the_class, polarity, :]
+        if polarity == 0:
+            true_positive_clause_outputs = clause_outputs[Y == the_class].sum(axis=0) / Y[Y == the_class].shape[0]
+        else:
+            true_positive_clause_outputs = clause_outputs[Y != the_class].sum(axis=0) / Y[Y != the_class].shape[0]
+        return true_positive_clause_outputs
+
+    def get_weight(self, the_class, polarity, clause):
+        if polarity == 0:
+            return self.weight_banks[the_class].get_weights()[clause]
+        else:
+            return self.weight_banks[the_class].get_weights()[self.number_of_clauses // 2 + clause]
+
+    def set_weight(self, the_class, polarity, clause, weight):
+        if polarity == 0:
+            self.weight_banks[the_class].get_weights()[clause] = weight
+        else:
+            self.weight_banks[the_class].get_weights()[self.number_of_clauses // 2 + clause] = weight
+
+    def get_ta_action(self, the_class, polarity, clause, ta):
+        if polarity == 0:
+            return self.clause_banks[the_class].get_ta_action(clause, ta)
+        else:
+            return self.clause_banks[the_class].get_ta_action(self.number_of_clauses // 2 + clause, ta)
+
+    def get_ta_state(self, the_class, polarity, clause, ta):
+        if polarity == 0:
+            return self.clause_banks[the_class].get_ta_state(clause, ta)
+        else:
+            return self.clause_banks[the_class].get_ta_state(self.number_of_clauses // 2 + clause, ta)
+
+    def set_ta_state(self, the_class, polarity, clause, ta, state):
+        if polarity == 0:
+            return self.clause_banks[the_class].set_ta_state(clause, ta, state)
+        else:
+            return self.clause_banks[the_class].set_ta_state(self.number_of_clauses // 2 + clause, ta, state)
+
+    def number_of_include_actions(self, the_class, clause):
+        return self.clause_banks[the_class].number_of_include_actions(clause)
diff --git a/tmu/models/regression/__init__.py b/tmu/models/regression/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tmu/models/regression/vanilla_regressor.py b/tmu/models/regression/vanilla_regressor.py
new file mode 100644
index 00000000..749c9875
--- /dev/null
+++ b/tmu/models/regression/vanilla_regressor.py
@@ -0,0 +1,117 @@
+# Copyright (c) 2023 Ole-Christoffer Granmo
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+from tmu.clause_bank import ClauseBank
+from tmu.models.base import TMBasis
+from tmu.weight_bank import WeightBank
+import numpy as np
+
+
+class TMRegressor(TMBasis):
+    def __init__(self, number_of_clauses, T, s, platform='CPU', patch_dim=None, feature_negation=True,
+                 boost_true_positive_feedback=1, max_included_literals=None, number_of_state_bits_ta=8,
+                 weighted_clauses=False, clause_drop_p=0.0, literal_drop_p=0.0):
+        super().__init__(number_of_clauses, T, s, platform=platform, patch_dim=patch_dim,
+                         feature_negation=feature_negation, boost_true_positive_feedback=boost_true_positive_feedback,
+                         max_included_literals=max_included_literals, number_of_state_bits_ta=number_of_state_bits_ta,
+                         weighted_clauses=weighted_clauses, clause_drop_p=clause_drop_p, literal_drop_p=literal_drop_p)
+
+    def initialize(self, X, Y):
+        self.max_y = np.max(Y)
+        self.min_y = np.min(Y)
+
+        if self.platform == 'CPU':
+            self.clause_bank = ClauseBank(X, self.number_of_clauses, self.number_of_state_bits_ta,
+                                          self.number_of_state_bits_ind, self.patch_dim)
+        elif self.platform == 'CUDA':
+            from tmu.clause_bank_cuda import ClauseBankCUDA
+            self.clause_bank = ClauseBankCUDA(X, self.number_of_clauses, self.number_of_state_bits_ta, self.patch_dim)
+        else:
+            print("Unknown Platform")
+            sys.exit(-1)
+
+        self.weight_bank = WeightBank(np.ones(self.number_of_clauses).astype(np.int32))
+
+        if self.max_included_literals == None:
+            self.max_included_literals = self.clause_bank.number_of_literals
+
+    def fit(self, X, Y, shuffle=True):
+        if self.initialized == False:
+            self.initialize(X, Y)
+            self.initialized = True
+
+        if not np.array_equal(self.X_train, X):
+            self.encoded_X_train = self.clause_bank.prepare_X(X)
+            self.X_train = X.copy()
+        encoded_Y = np.ascontiguousarray(((Y - self.min_y) / (self.max_y - self.min_y) * self.T).astype(np.int32))
+
+        clause_active = np.ascontiguousarray(
+            np.random.choice(2, self.number_of_clauses, p=[self.clause_drop_p, 1.0 - self.clause_drop_p]).astype(
+                np.int32))
+        literal_active = (np.zeros(self.clause_bank.number_of_ta_chunks, dtype=np.uint32) | ~0).astype(np.uint32)
+        if not self.feature_negation:
+            for k in range(self.clause_bank.number_of_literals // 2, self.clause_bank.number_of_literals):
+                ta_chunk = k // 32
+                chunk_pos = k % 32
+                literal_active[ta_chunk] &= (~(1 << chunk_pos))
+
+        literal_active = literal_active.astype(np.uint32)
+
+        shuffled_index = np.arange(X.shape[0])
+        if shuffle:
+            np.random.shuffle(shuffled_index)
+
+        for e in shuffled_index:
+            clause_outputs = self.clause_bank.calculate_clause_outputs_update(literal_active, self.encoded_X_train, e)
+
+            pred_y = np.dot(clause_active * self.weight_bank.get_weights(), clause_outputs).astype(np.int32)
+            pred_y = np.clip(pred_y, 0, self.T)
+            prediction_error = pred_y - encoded_Y[e];
+
+            update_p = (1.0 * prediction_error / self.T) ** 2
+
+            if pred_y < encoded_Y[e]:
+                self.clause_bank.type_i_feedback(update_p, self.s, self.boost_true_positive_feedback,
+                                                 self.max_included_literals, clause_active, literal_active,
+                                                 self.encoded_X_train, e)
+                if self.weighted_clauses:
+                    self.weight_bank.increment(clause_outputs, update_p, clause_active, False)
+            elif pred_y > encoded_Y[e]:
+                self.clause_bank.type_ii_feedback(update_p, clause_active, literal_active, self.encoded_X_train, e)
+                if self.weighted_clauses:
+                    self.weight_bank.decrement(clause_outputs, update_p, clause_active, False)
+        return
+
+    def predict(self, X):
+        if not np.array_equal(self.X_test, X):
+            self.encoded_X_test = self.clause_bank.prepare_X(X)
+            self.X_test = X.copy()
+
+        Y = np.ascontiguousarray(np.zeros(X.shape[0]))
+        for e in range(X.shape[0]):
+            clause_outputs = self.clause_bank.calculate_clause_outputs_predict(self.encoded_X_test, e)
+            pred_y = np.dot(self.weight_bank.get_weights(), clause_outputs).astype(np.int32)
+            Y[e] = 1.0 * pred_y * (self.max_y - self.min_y) / (self.T) + self.min_y
+        return Y
+
+    def get_weight(self, clause):
+        return self.weight_bank.get_weights()[clause]
+
+    def set_weight(self, clause, weight):
+        self.weight_banks.get_weights()[clause] = weight
diff --git a/tmu/tools.py b/tmu/tools.py
index a639e35f..982ae387 100644
--- a/tmu/tools.py
+++ b/tmu/tools.py
@@ -24,7 +24,7 @@
 import numpy as np
 from scipy.sparse import csr_matrix, csc_matrix
 
-from ._tools import ffi, lib
+from tmu.tmulib import ffi, lib
 
 def produce_autoencoder_examples(X_csr, X_csc, active_output, accumulation, number_of_ta_chunks, append_negated=True):
 	encoded_X = np.ascontiguousarray(np.empty(int(number_of_ta_chunks*active_output.shape[0]), dtype=np.uint32))
diff --git a/tmu/tsetlin_machine.py b/tmu/tsetlin_machine.py
deleted file mode 100644
index 493617e0..00000000
--- a/tmu/tsetlin_machine.py
+++ /dev/null
@@ -1,1342 +0,0 @@
-# Copyright (c) 2022 Ole-Christoffer Granmo
-
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-# SOFTWARE.
-
-# This code implements the Convolutional Tsetlin Machine from paper arXiv:1905.09688
-# https://arxiv.org/abs/1905.09688
-
-import sys
-import numpy as np
-from tmu.clause_bank import ClauseBank
-from tmu.weight_bank import WeightBank
-from scipy.sparse import csr_matrix, csc_matrix
-from sys import maxsize
-
-from time import time
-
-class TMBasis():
-	def __init__(self, number_of_clauses, T, s, confidence_driven_updating=False, type_i_ii_ratio = 1.0, type_iii_feedback=False, focused_negative_sampling=False, output_balancing=False, d=200.0, platform='CPU', patch_dim=None, feature_negation=True, boost_true_positive_feedback=1, max_included_literals=None, number_of_state_bits_ta=8, number_of_state_bits_ind=8, weighted_clauses=False, clause_drop_p = 0.0, literal_drop_p = 0.0, batch_size=100, incremental=True):
-		self.number_of_clauses = number_of_clauses
-		self.number_of_state_bits_ta = number_of_state_bits_ta
-		self.number_of_state_bits_ind = number_of_state_bits_ind
-		self.T = int(T)
-		self.s = s
-
-		self.confidence_driven_updating = confidence_driven_updating
-
-		if type_i_ii_ratio >= 1.0:
-			self.type_i_p = 1.0
-			self.type_ii_p = 1.0/type_i_ii_ratio
-		else:
-			self.type_i_p = type_i_ii_ratio
-			self.type_ii_p = 1.0
-
-		self.type_iii_feedback = type_iii_feedback
-		self.focused_negative_sampling = focused_negative_sampling
-		self.output_balancing = output_balancing
-		self.d = d
-		self.platform = platform
-		self.patch_dim = patch_dim
-		self.feature_negation = feature_negation
-		self.boost_true_positive_feedback = boost_true_positive_feedback
-		self.max_included_literals = max_included_literals
-		self.weighted_clauses = weighted_clauses
-
-		self.clause_drop_p = clause_drop_p
-		self.literal_drop_p = literal_drop_p
-
-		self.batch_size = batch_size
-
-		self.incremental = incremental
-
-		self.X_train = np.zeros(0, dtype=np.uint32)
-		self.X_test = np.zeros(0, dtype=np.uint32)
-
-		self.initialized = False
-
-	def clause_co_occurrence(self, X, percentage=False):
-		clause_outputs = csr_matrix(self.transform(X))
-		if percentage:
-			return clause_outputs.transpose().dot(clause_outputs).multiply(1.0/clause_outputs.sum(axis=0))
-		else:
-			return clause_outputs.transpose().dot(clause_outputs)
-
-	def transform(self, X):
-		encoded_X = self.clause_bank.prepare_X(X)
-		transformed_X = np.empty((X.shape[0], self.number_of_clauses), dtype=np.uint32)
-		for e in range(X.shape[0]):
-			transformed_X[e,:] = self.clause_bank.calculate_clause_outputs_predict(encoded_X, e)
-		return transformed_X
-
-	def transform_patchwise(self, X):
-		encoded_X = tmu.tools.encode(X, X.shape[0], self.number_of_patches, self.clause_bank.number_of_ta_chunks, self.dim, self.patch_dim, 0)
-		transformed_X = np.empty((X.shape[0], self.number_of_clauses*self.number_of_patches), dtype=np.uint32)
-		for e in range(X.shape[0]):
-			transformed_X[e,:] = self.clause_bank.calculate_clause_outputs_patchwise(encoded_X, e)
-		return transformed_X.reshape((X.shape[0], self.number_of_clauses, self.number_of_patches))
-
-	def literal_clause_frequency(self):
-		clause_active = np.ones(self.number_of_clauses, dtype=np.uint32)
-		return self.clause_bank.calculate_literal_clause_frequency(clause_active)
-	
-	def get_ta_action(self, clause, ta):
-		return self.clause_bank.get_ta_action(clause, ta)
-
-	def get_ta_state(self, clause, ta):
-		return self.clause_bank.get_ta_state(clause, ta)
-
-	def set_ta_state(self, clause, ta, state):
-		return self.clause_bank.set_ta_state(clause, ta, state)
-
-class TMClassifier(TMBasis):
-	def __init__(self, number_of_clauses, T, s, confidence_driven_updating=False, type_i_ii_ratio = 1.0, type_iii_feedback=False, d=200.0, platform='CPU', patch_dim=None, feature_negation=True, boost_true_positive_feedback=1, max_included_literals=None, number_of_state_bits_ta=8, number_of_state_bits_ind=8, weighted_clauses=False, clause_drop_p = 0.0, literal_drop_p = 0.0, batch_size=100, incremental=True):
-		super().__init__(number_of_clauses, T, s, confidence_driven_updating=confidence_driven_updating, type_i_ii_ratio = type_i_ii_ratio, type_iii_feedback=type_iii_feedback, d=d, platform=platform, patch_dim=patch_dim, feature_negation=feature_negation, boost_true_positive_feedback=boost_true_positive_feedback, max_included_literals=max_included_literals, number_of_state_bits_ta=number_of_state_bits_ta, number_of_state_bits_ind=number_of_state_bits_ind, weighted_clauses=weighted_clauses, clause_drop_p = clause_drop_p, literal_drop_p = literal_drop_p, batch_size=batch_size, incremental=incremental)
-
-	def initialize(self, X, Y):
-		self.number_of_classes = int(np.max(Y) + 1)
-
-		self.weight_banks = []
-		for i in range(self.number_of_classes):
-			self.weight_banks.append(WeightBank(np.concatenate((np.ones(self.number_of_clauses//2, dtype=np.int32), -1*np.ones(self.number_of_clauses//2, dtype=np.int32)))))
-		
-		self.clause_banks = []
-		if self.platform == 'CPU':
-			for i in range(self.number_of_classes):
-				self.clause_banks.append(ClauseBank(X, self.number_of_clauses, self.number_of_state_bits_ta, self.number_of_state_bits_ind, self.patch_dim, batch_size=self.batch_size, incremental=self.incremental))
-		elif self.platform == 'CUDA':
-			from tmu.clause_bank_cuda import ClauseBankCUDA
-			for i in range(self.number_of_classes):
-				self.clause_banks.append(ClauseBankCUDA(X, self.number_of_clauses, self.number_of_state_bits_ta, self.patch_dim))
-		else:
-			print("Unknown Platform")
-			sys.exit(-1)
-
-		if self.max_included_literals == None:
-			self.max_included_literals = self.clause_banks[0].number_of_literals
-
-		self.positive_clauses = np.concatenate((np.ones(self.number_of_clauses//2, dtype=np.int32), np.zeros(self.number_of_clauses//2, dtype=np.int32)))
-		self.negative_clauses = np.concatenate((np.zeros(self.number_of_clauses//2, dtype=np.int32), np.ones(self.number_of_clauses//2, dtype=np.int32)))
-
-	def fit(self, X, Y, shuffle=True):
-		if self.initialized == False:
-			self.initialize(X, Y)
-			self.initialized = True
-
-		if not np.array_equal(self.X_train, X):
-			self.encoded_X_train = self.clause_banks[0].prepare_X(X)
-			self.X_train = X.copy()
-
-		Ym = np.ascontiguousarray(Y).astype(np.uint32)
-
-		clause_active = []
-		for i in range(self.number_of_classes):
-		 	# Clauses are dropped based on their weights
-			class_clause_active = np.ascontiguousarray(np.ones(self.number_of_clauses, dtype=np.int32))
-			clause_score = np.abs(self.weight_banks[i].get_weights())
-			deactivate = np.random.choice(np.arange(self.number_of_clauses), size=int(self.number_of_clauses*self.clause_drop_p), p = clause_score / clause_score.sum())
-			for d in range(deactivate.shape[0]):
-				class_clause_active[deactivate[d]] = 0
-			clause_active.append(class_clause_active)
-
-		# Literals are dropped based on their frequency
-		literal_active = (np.zeros(self.clause_banks[0].number_of_ta_chunks, dtype=np.uint32) | ~0).astype(np.uint32)
-		literal_clause_frequency = self.literal_clause_frequency()
-		if literal_clause_frequency.sum() > 0:
-			deactivate = np.random.choice(np.arange(self.clause_banks[0].number_of_literals), size=int(self.clause_banks[0].number_of_literals*self.literal_drop_p), p = literal_clause_frequency / literal_clause_frequency.sum())
-		else:
-			deactivate = np.random.choice(np.arange(self.clause_banks[0].number_of_literals), size=int(self.clause_banks[0].number_of_literals*self.literal_drop_p))
-		for d in range(deactivate.shape[0]):
-			ta_chunk = deactivate[d] // 32
-			chunk_pos = deactivate[d] % 32
-			literal_active[ta_chunk] &= (~(1 << chunk_pos))
-
-		if not self.feature_negation:
-			for k in range(self.clause_banks[0].number_of_literals//2, self.clause_banks[0].number_of_literals):
-				ta_chunk = k // 32
-				chunk_pos = k % 32
-				literal_active[ta_chunk] &= (~(1 << chunk_pos))
-		literal_active = literal_active.astype(np.uint32)
-
-		shuffled_index = np.arange(X.shape[0])
-		if shuffle:
-			np.random.shuffle(shuffled_index)
-
-		for e in shuffled_index:
-			target = Ym[e]
-			
-			clause_outputs = self.clause_banks[target].calculate_clause_outputs_update(literal_active, self.encoded_X_train, e)
-			class_sum = np.dot(clause_active[target] * self.weight_banks[target].get_weights(), clause_outputs).astype(np.int32)
-			class_sum = np.clip(class_sum, -self.T, self.T)
-
-			if self.confidence_driven_updating:
-				update_p = 1.0*(self.T - np.absolute(class_sum))/self.T
-			else:
-				update_p = (self.T - class_sum)/(2*self.T)
-
-			if self.weighted_clauses:
-				self.weight_banks[target].increment(clause_outputs, update_p, clause_active[target], False)
-			self.clause_banks[target].type_i_feedback(update_p*self.type_i_p, self.s, self.boost_true_positive_feedback, self.max_included_literals, clause_active[target]*self.positive_clauses, literal_active, self.encoded_X_train, e)
-			self.clause_banks[target].type_ii_feedback(update_p*self.type_ii_p, clause_active[target]*self.negative_clauses, literal_active, self.encoded_X_train, e)
-			if self.type_iii_feedback:
-				self.clause_banks[target].type_iii_feedback(update_p, self.d, clause_active[target]*self.positive_clauses, literal_active, self.encoded_X_train, e, 1)
-				self.clause_banks[target].type_iii_feedback(update_p, self.d, clause_active[target]*self.negative_clauses, literal_active, self.encoded_X_train, e, 0)
-
-			not_target = np.random.randint(self.number_of_classes)
-			while not_target == target:
-				not_target = np.random.randint(self.number_of_classes)
-
-			clause_outputs = self.clause_banks[not_target].calculate_clause_outputs_update(literal_active, self.encoded_X_train, e)
-			class_sum = np.dot(clause_active[not_target] * self.weight_banks[not_target].get_weights(), clause_outputs).astype(np.int32)
-			class_sum = np.clip(class_sum, -self.T, self.T)
-
-			if self.confidence_driven_updating:
-				update_p = 1.0*(self.T - np.absolute(class_sum))/self.T
-			else:
-				update_p = (self.T + class_sum)/(2*self.T)
-		
-			if self.weighted_clauses:
-				self.weight_banks[not_target].decrement(clause_outputs, update_p, clause_active[not_target], False)			
-			self.clause_banks[not_target].type_i_feedback(update_p*self.type_i_p, self.s, self.boost_true_positive_feedback, self.max_included_literals, clause_active[not_target]*self.negative_clauses, literal_active, self.encoded_X_train, e)
-			self.clause_banks[not_target].type_ii_feedback(update_p*self.type_ii_p, clause_active[not_target]*self.positive_clauses, literal_active, self.encoded_X_train, e)
-			if self.type_iii_feedback:
-				self.clause_banks[not_target].type_iii_feedback(update_p, self.d, clause_active[not_target]*self.negative_clauses, literal_active, self.encoded_X_train, e, 1)
-				self.clause_banks[not_target].type_iii_feedback(update_p, self.d, clause_active[not_target]*self.positive_clauses, literal_active, self.encoded_X_train, e, 0)
-		return
-
-	def predict(self, X):
-		if not np.array_equal(self.X_test, X):
-			self.encoded_X_test = self.clause_banks[0].prepare_X(X)
-			self.X_test = X.copy()
-
-		Y = np.ascontiguousarray(np.zeros(X.shape[0], dtype=np.uint32))
-		for e in range(X.shape[0]):
-			max_class_sum = -self.T
-			max_class = 0
-			for i in range(self.number_of_classes):
-				class_sum = np.dot(self.weight_banks[i].get_weights(), self.clause_banks[i].calculate_clause_outputs_predict(self.encoded_X_test, e)).astype(np.int32)
-				class_sum = np.clip(class_sum, -self.T, self.T)
-				if class_sum > max_class_sum:
-					max_class_sum = class_sum
-					max_class = i
-			Y[e] = max_class
-		return Y
-
-	def transform(self, X):
-		encoded_X = self.clause_banks[0].prepare_X(X)
-		transformed_X = np.empty((X.shape[0], self.number_of_classes, self.number_of_clauses), dtype=np.uint32)
-		for e in range(X.shape[0]):
-			for i in range(self.number_of_classes):
-				transformed_X[e,i,:] = self.clause_banks[i].calculate_clause_outputs_predict(encoded_X, e)
-		return transformed_X.reshape((X.shape[0], self.number_of_classes*self.number_of_clauses))
-
-	def transform_patchwise(self, X):
-		encoded_X = tmu.tools.encode(X, X.shape[0], self.number_of_patches, self.number_of_ta_chunks, self.dim, self.patch_dim, 0)
-		transformed_X = np.empty((X.shape[0], self.number_of_classes, self.number_of_clauses//2*self.number_of_patches), dtype=np.uint32)
-		for e in range(X.shape[0]):
-			for i in range(self.number_of_classes):
-				transformed_X[e,i,:] = self.clause_bank[i].calculate_clause_outputs_patchwise(encoded_X, e)
-		return transformed_X.reshape((X.shape[0], self.number_of_classes*self.number_of_clauses, self.number_of_patches))
-
-	def literal_clause_frequency(self):
-		clause_active = np.ones(self.number_of_clauses, dtype=np.uint32)
-		literal_frequency = np.zeros(self.clause_banks[0].number_of_literals, dtype=np.uint32)
-		for i in range(self.number_of_classes):
-			literal_frequency += self.clause_banks[i].calculate_literal_clause_frequency(clause_active)
-		return literal_frequency
-
-	def literal_importance(self, the_class, negated_features=False, negative_polarity=False):
-		literal_frequency = np.zeros(self.clause_banks[0].number_of_literals, dtype=np.uint32)
-		if negated_features:
-			if negative_polarity:
-				literal_frequency[self.clause_banks[the_class].number_of_literals//2:] += self.clause_banks[the_class].calculate_literal_clause_frequency(self.negative_clauses)[self.clause_banks[the_class].number_of_literals//2:]
-			else:
-				literal_frequency[self.clause_banks[the_class].number_of_literals//2:] += self.clause_banks[the_class].calculate_literal_clause_frequency(self.positive_clauses)[self.clause_banks[the_class].number_of_literals//2:]
-		else:
-			if negative_polarity:
-				literal_frequency[:self.clause_banks[the_class].number_of_literals//2] += self.clause_banks[the_class].calculate_literal_clause_frequency(self.negative_clauses)[:self.clause_banks[the_class].number_of_literals//2]
-			else:
-				literal_frequency[:self.clause_banks[the_class].number_of_literals//2] += self.clause_banks[the_class].calculate_literal_clause_frequency(self.positive_clauses)[:self.clause_banks[the_class].number_of_literals//2]
-
-		return literal_frequency
-
-	def clause_precision(self, the_class, polarity, X, Y):
-		clause_outputs = self.transform(X).reshape(X.shape[0], self.number_of_classes, 2, self.number_of_clauses//2)[:,the_class, polarity,:]
-		if polarity == 0:
-			true_positive_clause_outputs = clause_outputs[Y==the_class].sum(axis=0)
-			false_positive_clause_outputs = clause_outputs[Y!=the_class].sum(axis=0)
-		else:
-			true_positive_clause_outputs = clause_outputs[Y!=the_class].sum(axis=0)
-			false_positive_clause_outputs = clause_outputs[Y==the_class].sum(axis=0)
-		return np.where(true_positive_clause_outputs + false_positive_clause_outputs == 0, 0, true_positive_clause_outputs/(true_positive_clause_outputs + false_positive_clause_outputs))
-
-	def clause_recall(self, the_class, polarity, X, Y):
-		clause_outputs = self.transform(X).reshape(X.shape[0], self.number_of_classes, 2, self.number_of_clauses//2)[:,the_class, polarity,:]
-		if polarity == 0:
-			true_positive_clause_outputs = clause_outputs[Y==the_class].sum(axis=0) / Y[Y==the_class].shape[0]
-		else:
-			true_positive_clause_outputs = clause_outputs[Y!=the_class].sum(axis=0) / Y[Y!=the_class].shape[0]
-		return true_positive_clause_outputs
-	
-	def get_weight(self, the_class, polarity, clause):
-		if polarity == 0:
-			return self.weight_banks[the_class].get_weights()[clause]
-		else:
-			return self.weight_banks[the_class].get_weights()[self.number_of_clauses//2 + clause]
-
-	def set_weight(self, the_class, polarity, clause, weight):
-		if polarity == 0:
-			self.weight_banks[the_class].get_weights()[clause] = weight
-		else:
-			self.weight_banks[the_class].get_weights()[self.number_of_clauses//2 + clause] = weight
-
-	def get_ta_action(self, the_class, polarity, clause, ta):
-		if polarity == 0:
-			return self.clause_banks[the_class].get_ta_action(clause, ta)
-		else:
-			return self.clause_banks[the_class].get_ta_action(self.number_of_clauses//2 + clause, ta)
-
-	def get_ta_state(self, the_class, polarity, clause, ta):
-		if polarity == 0:
-			return self.clause_banks[the_class].get_ta_state(clause, ta)
-		else:
-			return self.clause_banks[the_class].get_ta_state(self.number_of_clauses//2 + clause, ta)
-
-	def set_ta_state(self, the_class, polarity, clause, ta, state):
-		if polarity == 0:
-			return self.clause_banks[the_class].set_ta_state(clause, ta, state)
-		else:
-			return self.clause_banks[the_class].set_ta_state(self.number_of_clauses//2 + clause, ta, state)
-
-	def number_of_include_actions(self, the_class, clause):
-		return self.clause_banks[the_class].number_of_include_actions(clause)
-		
-class TMCoalescedClassifier(TMBasis):
-	def __init__(self, number_of_clauses, T, s, type_i_ii_ratio = 1.0, type_iii_feedback=False, focused_negative_sampling=False, output_balancing=False, d=200.0, platform = 'CPU', patch_dim=None, feature_negation=True, boost_true_positive_feedback=1, max_included_literals=None, number_of_state_bits_ta=8, number_of_state_bits_ind=8, weighted_clauses=False, clause_drop_p = 0.0, literal_drop_p = 0.0):
-		super().__init__(number_of_clauses, T, s, type_i_ii_ratio = type_i_ii_ratio, type_iii_feedback=type_iii_feedback, focused_negative_sampling=focused_negative_sampling, output_balancing=output_balancing, d=d, platform = platform, patch_dim=patch_dim, feature_negation=feature_negation, boost_true_positive_feedback=boost_true_positive_feedback, max_included_literals=max_included_literals, number_of_state_bits_ta=number_of_state_bits_ta, number_of_state_bits_ind=number_of_state_bits_ind, weighted_clauses=weighted_clauses, clause_drop_p = clause_drop_p, literal_drop_p = literal_drop_p)
-
-	def initialize(self, X, Y):
-		self.number_of_classes = int(np.max(Y) + 1)
-	
-		if self.platform == 'CPU':
-			self.clause_bank = ClauseBank(X, self.number_of_clauses, self.number_of_state_bits_ta, self.number_of_state_bits_ind, self.patch_dim)
-		elif self.platform == 'CUDA':
-			from tmu.clause_bank_cuda import ClauseBankCUDA
-			self.clause_bank = ClauseBankCUDA(X, self.number_of_clauses, self.number_of_state_bits_ta, self.patch_dim)
-		else:
-			print("Unknown Platform")
-			sys.exit(-1)
-
-		self.weight_banks = []
-		for i in range(self.number_of_classes):
-			self.weight_banks.append(WeightBank(np.random.choice([-1,1], size=self.number_of_clauses).astype(np.int32)))
-	
-		if self.max_included_literals == None:
-			self.max_included_literals = self.clause_bank.number_of_literals
-
-	def update(self, target, e):
-		clause_outputs = self.clause_bank.calculate_clause_outputs_update(self.literal_active, self.encoded_X_train, e)
-			
-		class_sum = np.dot(self.clause_active * self.weight_banks[target].get_weights(), clause_outputs).astype(np.int32)
-		class_sum = np.clip(class_sum, -self.T, self.T)
-		update_p = (self.T - class_sum)/(2*self.T)
-
-		type_iii_feedback_selection = np.random.choice(2)
-
-		self.clause_bank.type_i_feedback(update_p*self.type_i_p, self.s, self.boost_true_positive_feedback, self.max_included_literals, self.clause_active*(self.weight_banks[target].get_weights() >= 0), self.literal_active, self.encoded_X_train, e)
-		self.clause_bank.type_ii_feedback(update_p*self.type_ii_p, self.clause_active*(self.weight_banks[target].get_weights() < 0), self.literal_active, self.encoded_X_train, e)
-		self.weight_banks[target].increment(clause_outputs, update_p, self.clause_active, True)
-		if self.type_iii_feedback and type_iii_feedback_selection == 0:
-			self.clause_bank.type_iii_feedback(update_p, self.d, self.clause_active*(self.weight_banks[target].get_weights() >= 0), self.literal_active, self.encoded_X_train, e, 1)
-			self.clause_bank.type_iii_feedback(update_p, self.d, self.clause_active*(self.weight_banks[target].get_weights() < 0), self.literal_active, self.encoded_X_train, e, 0)
-
-		for i in range(self.number_of_classes):
-			if i == target:
-				self.update_ps[i] = 0.0
-			else:
-				self.update_ps[i] = np.dot(self.clause_active * self.weight_banks[i].get_weights(), clause_outputs).astype(np.int32)
-				self.update_ps[i] = np.clip(self.update_ps[i], -self.T, self.T)
-				self.update_ps[i] = 1.0*(self.T + self.update_ps[i])/(2*self.T)
-
-		if self.update_ps.sum() == 0:
-			return
-
-		if self.focused_negative_sampling:
-			not_target = np.random.choice(self.number_of_classes, p=self.update_ps/self.update_ps.sum())
-			update_p = self.update_ps[not_target] 
-		else: 
-			not_target = np.random.randint(self.number_of_classes)
-			while not_target == target:
-				not_target = np.random.randint(self.number_of_classes)
-			update_p = self.update_ps[not_target]
-	
-		self.clause_bank.type_i_feedback(update_p*self.type_i_p, self.s, self.boost_true_positive_feedback, self.max_included_literals, self.clause_active * (self.weight_banks[not_target].get_weights() < 0), self.literal_active, self.encoded_X_train, e)
-		self.clause_bank.type_ii_feedback(update_p*self.type_ii_p, self.clause_active*(self.weight_banks[not_target].get_weights() >= 0), self.literal_active, self.encoded_X_train, e)
-		if self.type_iii_feedback and type_iii_feedback_selection == 1:
-			self.clause_bank.type_iii_feedback(update_p, self.d, self.clause_active*(self.weight_banks[not_target].get_weights() < 0), self.literal_active, self.encoded_X_train, e, 1)
-			self.clause_bank.type_iii_feedback(update_p, self.d, self.clause_active*(self.weight_banks[not_target].get_weights() >= 0), self.literal_active, self.encoded_X_train, e, 0)
-
-		self.weight_banks[not_target].decrement(clause_outputs, update_p, self.clause_active, True)
-
-	def fit(self, X, Y, shuffle=True):
-		if self.initialized == False:
-			self.initialize(X, Y)
-			self.initialized = True
-
-		if not np.array_equal(self.X_train, X):
-			self.encoded_X_train = self.clause_bank.prepare_X(X)
-			self.X_train = X.copy()
-
-		Ym = np.ascontiguousarray(Y).astype(np.uint32)
-
-		# Clauses are dropped based on their weights
-		self.clause_active = np.ones(self.number_of_clauses, dtype=np.uint32)
-		clause_score = np.zeros(self.number_of_clauses, dtype=np.int32)
-		for i in range(self.number_of_classes):
-			clause_score += np.abs(self.weight_banks[i].get_weights())
-		deactivate = np.random.choice(np.arange(self.number_of_clauses), size=int(self.number_of_clauses*self.clause_drop_p), p = clause_score / clause_score.sum())
-		for d in range(deactivate.shape[0]):
-			self.clause_active[deactivate[d]] = 0
-
-		# Literals are dropped based on their frequency
-		self.literal_active = (np.zeros(self.clause_bank.number_of_ta_chunks, dtype=np.uint32) | ~0).astype(np.uint32)
-		literal_clause_frequency = self.literal_clause_frequency()
-		if literal_clause_frequency.sum() > 0:
-			deactivate = np.random.choice(np.arange(self.clause_bank.number_of_literals), size=int(self.clause_bank.number_of_literals*self.literal_drop_p), p = literal_clause_frequency / literal_clause_frequency.sum())
-		else:
-			deactivate = np.random.choice(np.arange(self.clause_bank.number_of_literals), size=int(self.clause_bank.number_of_literals*self.literal_drop_p))
-		for d in range(deactivate.shape[0]):
-			ta_chunk = deactivate[d] // 32
-			chunk_pos = deactivate[d] % 32
-			self.literal_active[ta_chunk] &= (~(1 << chunk_pos))
-
-		if not self.feature_negation:
-			for k in range(self.clause_bank.number_of_literals//2, self.clause_bank.number_of_literals):
-				ta_chunk = k // 32
-				chunk_pos = k % 32
-				self.literal_active[ta_chunk] &= (~(1 << chunk_pos))
-
-		self.literal_active = self.literal_active.astype(np.uint32)
-
-		self.update_ps = np.empty(self.number_of_classes)
-
-		shuffled_index = np.arange(X.shape[0])
-		if shuffle:
-			np.random.shuffle(shuffled_index)
-
-		class_observed = np.zeros(self.number_of_classes, dtype=np.uint32)
-		example_indexes = np.zeros(self.number_of_classes, dtype=np.uint32)
-		example_counter = 0
-		for e in shuffled_index:
-			if self.output_balancing:
-				if class_observed[Ym[e]] == 0:
-					example_indexes[Ym[e]] = e
-					class_observed[Ym[e]] = 1
-					example_counter += 1
-			else:
-				example_indexes[example_counter] = e
-				example_counter += 1
-
-			if example_counter == self.number_of_classes:
-				example_counter = 0
-
-				for i in range(self.number_of_classes):
-					class_observed[i] = 0
-					batch_example = example_indexes[i]
-					self.update(Ym[batch_example], batch_example)
-		return
-
-	def predict(self, X):
-		if not np.array_equal(self.X_test, X):
-			self.encoded_X_test = self.clause_bank.prepare_X(X)
-			self.X_test = X.copy()
-
-		Y = np.ascontiguousarray(np.zeros(X.shape[0], dtype=np.uint32))
-
-		for e in range(X.shape[0]):
-			max_class_sum = -self.T
-			max_class = 0
-			clause_outputs = self.clause_bank.calculate_clause_outputs_predict(self.encoded_X_test, e)			
-			for i in range(self.number_of_classes):
-				class_sum = np.dot(self.weight_banks[i].get_weights(), clause_outputs).astype(np.int32)
-				class_sum = np.clip(class_sum, -self.T, self.T)
-				if class_sum > max_class_sum:
-					max_class_sum = class_sum
-					max_class = i
-			Y[e] = max_class
-		return Y
-
-	def predict_individual(self, X):
-		if not np.array_equal(self.X_test, X):
-			self.encoded_X_test = self.clause_bank.prepare_X(X)
-			self.X_test = X.copy()
-
-		Y = np.ascontiguousarray(np.zeros((X.shape[0], self.number_of_classes), dtype=np.uint32))
-
-		for e in range(X.shape[0]):
-			clause_outputs = self.clause_bank.calculate_clause_outputs_predict(self.encoded_X_test, e)			
-			for i in range(self.number_of_classes):
-				class_sum = np.dot(self.weight_banks[i].get_weights(), clause_outputs).astype(np.int32)
-				class_sum = np.clip(class_sum, -self.T, self.T)
-				Y[e, i] = (class_sum >= 0)
-		return Y
-
-	def clause_precision(self, the_class, positive_polarity, X, Y):
-		clause_outputs = self.transform(X)
-		weights = self.weight_banks[the_class].get_weights()
-		if positive_polarity == 0:
-			positive_clause_outputs = (weights >= 0)[:,np.newaxis].transpose() * clause_outputs
-			true_positive_clause_outputs = positive_clause_outputs[Y==the_class].sum(axis=0)
-			false_positive_clause_outputs = positive_clause_outputs[Y!=the_class].sum(axis=0)
-		else:
-			positive_clause_outputs = (weights < 0)[:,np.newaxis].transpose() * clause_outputs
-			true_positive_clause_outputs = positive_clause_outputs[Y!=the_class].sum(axis=0)
-			false_positive_clause_outputs = positive_clause_outputs[Y==the_class].sum(axis=0)
-		
-		return np.where(true_positive_clause_outputs + false_positive_clause_outputs == 0, 0, 1.0*true_positive_clause_outputs/(true_positive_clause_outputs + false_positive_clause_outputs))
-
-	def clause_recall(self, the_class, positive_polarity, X, Y):
-		clause_outputs = self.transform(X)
-		weights = self.weight_banks[the_class].get_weights()
-		
-		if positive_polarity == 0:
-			positive_clause_outputs = (weights >= 0)[:,np.newaxis].transpose() * clause_outputs
-			true_positive_clause_outputs = positive_clause_outputs[Y==the_class].sum(axis=0) / Y[Y==the_class].shape[0]
-		else:
-			positive_clause_outputs = (weights < 0)[:,np.newaxis].transpose() * clause_outputs
-			true_positive_clause_outputs = positive_clause_outputs[Y!=the_class].sum(axis=0) / Y[Y!=the_class].shape[0]
-			
-		return true_positive_clause_outputs 
-
-	def get_weight(self, the_class, clause):
-		return self.weight_banks[the_class].get_weights()[clause]
-
-	def set_weight(self, the_class, clause, weight):
-		self.weight_banks[the_class].get_weights()[clause] = weight
-
-	def number_of_include_actions(self, clause):
-		return self.clause_bank.number_of_include_actions(clause)
-
-class TMAutoEncoder(TMBasis):
-	def __init__(self, number_of_clauses, T, s, output_active, accumulation=1, type_i_ii_ratio = 1.0, type_iii_feedback=False, focused_negative_sampling=False, output_balancing=False, d=200.0, platform = 'CPU', patch_dim=None, feature_negation=True, boost_true_positive_feedback=1, max_included_literals=None, number_of_state_bits_ta=8, number_of_state_bits_ind=8, weighted_clauses=False, clause_drop_p = 0.0, literal_drop_p = 0.0):
-		self.output_active = output_active
-		self.accumulation = accumulation
-		super().__init__(number_of_clauses, T, s, type_i_ii_ratio = type_i_ii_ratio, type_iii_feedback=type_iii_feedback, focused_negative_sampling=focused_negative_sampling, output_balancing=output_balancing, d=d, platform = platform, patch_dim=patch_dim, feature_negation=feature_negation, boost_true_positive_feedback=boost_true_positive_feedback, max_included_literals=max_included_literals, number_of_state_bits_ta=number_of_state_bits_ta, number_of_state_bits_ind=number_of_state_bits_ind, weighted_clauses=weighted_clauses, clause_drop_p = clause_drop_p, literal_drop_p = literal_drop_p)
-
-	def initialize(self, X):
-		self.number_of_classes = self.output_active.shape[0]
-		if self.platform == 'CPU':
-			self.clause_bank = ClauseBank(X, self.number_of_clauses, self.number_of_state_bits_ta, self.number_of_state_bits_ind, self.patch_dim)
-		elif self.platform == 'CUDA':
-			from tmu.clause_bank_cuda import ClauseBankCUDA
-			self.clause_bank = ClauseBankCUDA(X, self.number_of_clauses, self.number_of_state_bits_ta, self.patch_dim)
-		else:
-			print("Unknown Platform")
-			sys.exit(-1)
-
-		self.weight_banks = []
-		for i in range(self.number_of_classes):
-			self.weight_banks.append(WeightBank(np.random.choice([-1,1], size=self.number_of_clauses).astype(np.int32)))
-	
-		if self.max_included_literals == None:
-			self.max_included_literals = self.clause_bank.number_of_literals
-
-	def update(self, target_output, target_value, encoded_X, clause_active, literal_active):
-		all_literal_active = (np.zeros(self.clause_bank.number_of_ta_chunks, dtype=np.uint32) | ~0).astype(np.uint32)
-		clause_outputs = self.clause_bank.calculate_clause_outputs_update(all_literal_active, encoded_X, 0)
-			
-		class_sum = np.dot(clause_active * self.weight_banks[target_output].get_weights(), clause_outputs).astype(np.int32)
-		class_sum = np.clip(class_sum, -self.T, self.T)
-		
-		type_iii_feedback_selection = np.random.choice(2)
-
-		if target_value == 1:
-			update_p = (self.T - class_sum)/(2*self.T)
-
-			self.clause_bank.type_i_feedback(update_p*self.type_i_p, self.s, self.boost_true_positive_feedback, self.max_included_literals, clause_active*(self.weight_banks[target_output].get_weights() >= 0), literal_active, encoded_X, 0)
-			self.clause_bank.type_ii_feedback(update_p*self.type_ii_p, clause_active*(self.weight_banks[target_output].get_weights() < 0), literal_active, encoded_X, 0)
-			self.weight_banks[target_output].increment(clause_outputs, update_p, clause_active, True)
-			if self.type_iii_feedback and type_iii_feedback_selection == 0:
-				self.clause_bank.type_iii_feedback(update_p, self.d, clause_active*(self.weight_banks[target_output].get_weights() >= 0), literal_active, encoded_X, 0, 1)
-				self.clause_bank.type_iii_feedback(update_p, self.d, clause_active*(self.weight_banks[target_output].get_weights() < 0), literal_active, encoded_X, 0, 0)
-		else:
-			update_p = (self.T + class_sum)/(2*self.T)
-
-			self.clause_bank.type_i_feedback(update_p*self.type_i_p, self.s, self.boost_true_positive_feedback, self.max_included_literals, clause_active * (self.weight_banks[target_output].get_weights() < 0), literal_active, encoded_X, 0)
-			self.clause_bank.type_ii_feedback(update_p*self.type_ii_p, clause_active*(self.weight_banks[target_output].get_weights() >= 0), literal_active, encoded_X, 0)
-			self.weight_banks[target_output].decrement(clause_outputs, update_p, clause_active, True)
-			if self.type_iii_feedback and type_iii_feedback_selection == 1:
-				self.clause_bank.type_iii_feedback(update_p, self.d, clause_active*(self.weight_banks[target_output].get_weights() < 0), literal_active, encoded_X, 0, 1)
-				self.clause_bank.type_iii_feedback(update_p, self.d, clause_active*(self.weight_banks[target_output].get_weights() >= 0), literal_active, encoded_X, 0, 0)
-		return
-
-	def activate_clauses(self):
-		# Clauses are dropped based on their weights
-		clause_active = np.ones(self.number_of_clauses, dtype=np.uint32)
-		deactivate = np.random.choice(self.number_of_clauses, size=int(self.number_of_clauses*self.clause_drop_p))
-		for d in range(deactivate.shape[0]):
-			clause_active[deactivate[d]] = 0
-
-		return clause_active
-
-	def activate_literals(self):		
-		# Literals are dropped based on their frequency
-		literal_active = (np.zeros(self.clause_bank.number_of_ta_chunks, dtype=np.uint32) | ~0).astype(np.uint32)
-		literal_clause_frequency = self.literal_clause_frequency()
-		deactivate = np.random.choice(self.clause_bank.number_of_literals, size=int(self.clause_bank.number_of_literals*self.literal_drop_p))
-		for d in range(deactivate.shape[0]):
-			ta_chunk = deactivate[d] // 32
-			chunk_pos = deactivate[d] % 32
-			literal_active[ta_chunk] &= (~(1 << chunk_pos))
-
-		if not self.feature_negation:
-			for k in range(self.clause_bank.number_of_literals//2, self.clause_bank.number_of_literals):
-				ta_chunk = k // 32
-				chunk_pos = k % 32
-				literal_active[ta_chunk] &= (~(1 << chunk_pos))
-		literal_active = literal_active.astype(np.uint32)
-
-		return(literal_active)
-
-	def fit(self, X, number_of_examples=2000, shuffle=True):
-		if self.initialized == False:
-			self.initialize(X)
-			self.initialized = True
-
-		X_csr = csr_matrix(X.reshape(X.shape[0], -1))
-		X_csc = csc_matrix(X.reshape(X.shape[0], -1)).sorted_indices()
-
-		clause_active = self.activate_clauses()
-		literal_active = self.activate_literals()
-		
-		class_index = np.arange(self.number_of_classes, dtype=np.uint32)
-		for e in range(number_of_examples):
-			np.random.shuffle(class_index)
-
-			average_absolute_weights = np.zeros(self.number_of_clauses, dtype=np.float32)
-			for i in class_index:
-				 average_absolute_weights += np.absolute(self.weight_banks[i].get_weights())
-			average_absolute_weights /= self.number_of_classes
-			update_clause = np.random.random(self.number_of_clauses) <= (self.T - np.clip(average_absolute_weights, 0, self.T))/self.T
-
-			Xu, Yu = self.clause_bank.prepare_autoencoder_examples(X_csr, X_csc, self.output_active, self.accumulation)
-			for i in class_index:
-				(target, encoded_X) = Yu[i], Xu[i].reshape((1,-1))
-
-				ta_chunk = self.output_active[i] // 32
-				chunk_pos = self.output_active[i] % 32
-				copy_literal_active_ta_chunk = literal_active[ta_chunk]
-				literal_active[ta_chunk] &= ~(1 << chunk_pos)
-
-				self.update(i, target, encoded_X, update_clause*clause_active, literal_active)
-				literal_active[ta_chunk] = copy_literal_active_ta_chunk
-		return
-
-	def predict(self, X):
-		X_csr = csr_matrix(X.reshape(X.shape[0], -1))
-		Y = np.ascontiguousarray(np.zeros((self.number_of_classes, X.shape[0]), dtype=np.uint32))
-
-		for e in range(X.shape[0]):
-			encoded_X = self.clause_bank.prepare_X(X_csr[e,:].toarray())		
-
-			clause_outputs = self.clause_bank.calculate_clause_outputs_predict(encoded_X, 0)[0]		
-			for i in range(self.number_of_classes):
-				class_sum = np.dot(self.weight_banks[i].get_weights(), clause_outputs).astype(np.int32)
-				Y[i, e] = (class_sum >= 0)
-		return Y
-
-	def literal_importance(self, the_class, negated_features=False, negative_polarity=False):
-		literal_frequency = np.zeros(self.clause_bank.number_of_literals, dtype=np.uint32)
-		if negated_features:
-			if negative_polarity:
-				literal_frequency[self.clause_bank.number_of_literals//2:] += self.clause_bank.calculate_literal_clause_frequency(self.weight_banks[the_class].get_weights() < 0)[self.clause_bank.number_of_literals//2:]
-			else:
-				literal_frequency[self.clause_bank.number_of_literals//2:] += self.clause_bank.calculate_literal_clause_frequency(self.weight_banks[the_class].get_weights() >= 0)[self.clause_bank.number_of_literals//2:]
-		else:
-			if negative_polarity:
-				literal_frequency[:self.clause_bank.number_of_literals//2] += self.clause_bank.calculate_literal_clause_frequency(self.weight_banks[the_class].get_weights() < 0)[:self.clause_bank.number_of_literals//2]
-			else:
-				literal_frequency[:self.clause_bank.number_of_literals//2] += self.clause_bank.calculate_literal_clause_frequency(self.weight_banks[the_class].get_weights() >= 0)[:self.clause_bank.number_of_literals//2]
-
-		return literal_frequency
-
-	def clause_precision(self, the_class, positive_polarity, X, number_of_examples=2000):
-		X_csr = csr_matrix(X.reshape(X.shape[0], -1))
-		X_csc = csc_matrix(X.reshape(X.shape[0], -1)).sorted_indices()
-
-		true_positive = np.zeros(self.number_of_clauses, dtype=np.uint32)
-		false_positive = np.zeros(self.number_of_clauses, dtype=np.uint32)
-
-		weights = self.weight_banks[the_class].get_weights()
-
-		clause_active = self.activate_clauses()
-		literal_active = self.activate_literals()
-
-		for e in range(number_of_examples):
-			Xu, Yu = self.clause_bank.prepare_autoencoder_examples(X_csr, X_csc, np.array([self.output_active[the_class]], dtype=np.uint32), self.accumulation)
-			(target, encoded_X) = Yu[0], Xu[0].reshape((1,-1))
-			
-			clause_outputs = self.clause_bank.calculate_clause_outputs_predict(encoded_X, 0)
-
-			if positive_polarity:
-				if target == 1:
-					true_positive += (weights >= 0) * clause_outputs
-				else:
-					false_positive += (weights >= 0) * clause_outputs
-			else:
-				if target == 0:
-					true_positive += (weights < 0) * clause_outputs
-				else:
-					false_positive += (weights < 0) * clause_outputs				
-
-		return 1.0*true_positive/(true_positive + false_positive)
-
-	def clause_recall(self, the_class, positive_polarity, X, number_of_examples=2000):
-		X_csr = csr_matrix(X.reshape(X.shape[0], -1))
-		X_csc = csc_matrix(X.reshape(X.shape[0], -1)).sorted_indices()
-		
-		true_positive = np.zeros(self.number_of_clauses, dtype=np.uint32)
-		false_negative = np.zeros(self.number_of_clauses, dtype=np.uint32)
-
-		weights = self.weight_banks[the_class].get_weights()
-
-		for e in range(number_of_examples):
-			Xu, Yu = self.clause_bank.prepare_autoencoder_examples(X_csr, X_csc, np.array([self.output_active[the_class]], dtype=np.uint32), self.accumulation)
-			(target, encoded_X) = Yu[0], Xu[0].reshape((1,-1))
-
-			clause_outputs = self.clause_bank.calculate_clause_outputs_predict(encoded_X, 0)
-
-			if positive_polarity:
-				if target == 1:
-					true_positive += (weights >= 0) * clause_outputs
-					false_negative += (weights >= 0) * (1 - clause_outputs)
-			else:
-				if target == 0:
-					true_positive += (weights < 0) * clause_outputs
-					false_negative += (weights < 0) * (1 - clause_outputs)
-
-		return true_positive/(true_positive + false_negative)
-
-	def get_weight(self, the_class, clause):
-		return self.weight_banks[the_class].get_weights()[clause]
-
-	def get_weights(self, the_class):
-		return self.weight_banks[the_class].get_weights()
-
-	def set_weight(self, the_class, clause, weight):
-		self.weight_banks[the_class].get_weights()[clause] = weight
-
-class TMMultiTaskClassifier(TMBasis):
-	def __init__(self, number_of_clauses, T, s, confidence_driven_updating=False, type_i_ii_ratio = 1.0, type_iii_feedback=False, focused_negative_sampling=False, output_balancing=False, d=200.0, platform = 'CPU', patch_dim=None, feature_negation=True, boost_true_positive_feedback=1, max_included_literals=None, number_of_state_bits_ta=8, number_of_state_bits_ind=8, weighted_clauses=False, clause_drop_p = 0.0, literal_drop_p = 0.0):
-		super().__init__(number_of_clauses, T, s, confidence_driven_updating=confidence_driven_updating, type_i_ii_ratio = type_i_ii_ratio, type_iii_feedback=type_iii_feedback, focused_negative_sampling=focused_negative_sampling, output_balancing=output_balancing, d=d, platform = platform, patch_dim=patch_dim, feature_negation=feature_negation, boost_true_positive_feedback=boost_true_positive_feedback, max_included_literals=max_included_literals, number_of_state_bits_ta=number_of_state_bits_ta, number_of_state_bits_ind=number_of_state_bits_ind, weighted_clauses=weighted_clauses, clause_drop_p = clause_drop_p, literal_drop_p = literal_drop_p)
-
-	def initialize(self, X, Y):
-		self.number_of_classes = len(X)
-	
-		if self.platform == 'CPU':
-			self.clause_bank = ClauseBank(X[0], self.number_of_clauses, self.number_of_state_bits_ta, self.number_of_state_bits_ind, self.patch_dim)
-		elif self.platform == 'CUDA':
-			from tmu.clause_bank_cuda import ClauseBankCUDA
-			self.clause_bank = ClauseBankCUDA(X[0], self.number_of_clauses, self.number_of_state_bits_ta, self.patch_dim)
-		else:
-			print("Unknown Platform")
-			sys.exit(-1)
-
-		self.weight_banks = []
-		for i in range(self.number_of_classes):
-			self.weight_banks.append(WeightBank(np.random.choice([-1,1], size=self.number_of_clauses).astype(np.int32)))
-		
-		if self.max_included_literals == None:
-			self.max_included_literals = self.clause_bank.number_of_literals
-
-	def fit(self, X, Y, shuffle=True):
-		if self.initialized == False:
-			self.initialize(X, Y)
-			self.initialized = True
-
-		X_csr = {}
-		for i in range(self.number_of_classes):
-			X_csr[i] = csr_matrix(X[i].reshape(X[i].shape[0], -1))
-
-		# Clauses are dropped based on their weights
-		self.clause_active = np.ones(self.number_of_clauses, dtype=np.uint32)
-		clause_score = np.zeros(self.number_of_clauses, dtype=np.int32)
-		for i in range(self.number_of_classes):
-			clause_score += np.abs(self.weight_banks[i].get_weights())
-		deactivate = np.random.choice(np.arange(self.number_of_clauses), size=int(self.number_of_clauses*self.clause_drop_p), p = clause_score / clause_score.sum())
-		for d in range(deactivate.shape[0]):
-			self.clause_active[deactivate[d]] = 0
-
-		# Literals are dropped based on their frequency
-		self.literal_active = (np.zeros(self.clause_bank.number_of_ta_chunks, dtype=np.uint32) | ~0).astype(np.uint32)
-		literal_clause_frequency = self.literal_clause_frequency()
-		if literal_clause_frequency.sum() > 0:
-			deactivate = np.random.choice(np.arange(self.clause_bank.number_of_literals), size=int(self.clause_bank.number_of_literals*self.literal_drop_p), p = literal_clause_frequency / literal_clause_frequency.sum())
-		else:
-			deactivate = np.random.choice(np.arange(self.clause_bank.number_of_literals), size=int(self.clause_bank.number_of_literals*self.literal_drop_p))
-		for d in range(deactivate.shape[0]):
-			ta_chunk = deactivate[d] // 32
-			chunk_pos = deactivate[d] % 32
-			self.literal_active[ta_chunk] &= (~(1 << chunk_pos))
-
-		if not self.feature_negation:
-			for k in range(self.clause_bank.number_of_literals//2, self.clause_bank.number_of_literals):
-				ta_chunk = k // 32
-				chunk_pos = k % 32
-				self.literal_active[ta_chunk] &= (~(1 << chunk_pos))
-
-		self.literal_active = self.literal_active.astype(np.uint32)
-
-		shuffled_index = np.arange(X[0].shape[0])
-		if shuffle:
-			np.random.shuffle(shuffled_index)
-
-		class_index = np.arange(self.number_of_classes, dtype=np.uint32)
-		for e in shuffled_index:
-			np.random.shuffle(class_index)
-
-			average_absolute_weights = np.zeros(self.number_of_clauses, dtype=np.float32)
-			for i in class_index:
-				 average_absolute_weights += np.absolute(self.weight_banks[i].get_weights())
-			average_absolute_weights /= self.number_of_classes
-			update_clause = np.random.random(self.number_of_clauses) <= (self.T - np.clip(average_absolute_weights, 0, self.T))/self.T
-
-			for i in class_index:
-				encoded_X = self.clause_bank.prepare_X(X_csr[i][e,:].toarray())
-				clause_outputs = self.clause_bank.calculate_clause_outputs_update(self.literal_active, encoded_X, 0)
-			
-				class_sum = np.dot(self.clause_active * self.weight_banks[i].get_weights(), clause_outputs).astype(np.int32)
-				class_sum = np.clip(class_sum, -self.T, self.T)
-
-				type_iii_feedback_selection = np.random.choice(2)
-
-				if Y[i][e] == 1:
-					if self.confidence_driven_updating:
-						update_p = 1.0*(self.T - np.absolute(class_sum))/self.T
-					else:
-						update_p = (self.T - class_sum)/(2*self.T)
-
-					self.clause_bank.type_i_feedback(update_p*self.type_i_p, self.s, self.boost_true_positive_feedback, self.max_included_literals, update_clause*self.clause_active*(self.weight_banks[i].get_weights() >= 0), self.literal_active, encoded_X, 0)
-					self.clause_bank.type_ii_feedback(update_p*self.type_ii_p, update_clause*self.clause_active*(self.weight_banks[i].get_weights() < 0), self.literal_active, encoded_X, 0)
-					self.weight_banks[i].increment(clause_outputs, update_p, update_clause*self.clause_active, True)
-					if self.type_iii_feedback and type_iii_feedback_selection == 0:
-						self.clause_bank.type_iii_feedback(update_p, self.d, update_clause*self.clause_active*(self.weight_banks[i].get_weights() >= 0), self.literal_active, encoded_X, 0, 1)
-						self.clause_bank.type_iii_feedback(update_p, self.d, update_clause*self.clause_active*(self.weight_banks[i].get_weights() < 0), self.literal_active, encoded_X, 0, 0)
-				else:
-					if self.confidence_driven_updating:
-						update_p = 1.0*(self.T - np.absolute(class_sum))/self.T
-					else:
-						update_p = (self.T + class_sum)/(2*self.T)
-
-					self.clause_bank.type_i_feedback(update_p*self.type_i_p, self.s, self.boost_true_positive_feedback, self.max_included_literals, update_clause*self.clause_active*(self.weight_banks[i].get_weights() < 0), self.literal_active, encoded_X, 0)
-					self.clause_bank.type_ii_feedback(update_p*self.type_ii_p, update_clause*update_clause*self.clause_active*(self.weight_banks[i].get_weights() >= 0), self.literal_active, encoded_X, 0)
-					self.weight_banks[i].decrement(clause_outputs, update_p, self.clause_active, True)
-					if self.type_iii_feedback and type_iii_feedback_selection == 1:
-						self.clause_bank.type_iii_feedback(update_p, self.d, update_clause*self.clause_active*(self.weight_banks[i].get_weights() < 0), self.literal_active, encoded_X, 0, 1)
-						self.clause_bank.type_iii_feedback(update_p, self.d, update_clause*self.clause_active*(self.weight_banks[i].get_weights() >= 0), self.literal_active, encoded_X, 0, 0)
-		return
-
-	def predict(self, X):
-		Y = {}
-		for i in range(len(X)):
-			Y[i] = np.ascontiguousarray(np.zeros((X[0].shape[0]), dtype=np.uint32))
-
-		X_csr = {}
-		for i in range(self.number_of_classes):
-			X_csr[i] = csr_matrix(X[i].reshape(X[i].shape[0], -1))
-		
-		for i in range(self.number_of_classes):
-			for e in range(X[i].shape[0]):
-				encoded_X = self.clause_bank.prepare_X(X_csr[i][e,:].toarray())		
-				clause_outputs = self.clause_bank.calculate_clause_outputs_predict(encoded_X, 0)			
-				class_sum = np.dot(self.weight_banks[i].get_weights(), clause_outputs).astype(np.int32)
-				Y[i][e] = (class_sum >= 0)
-		return Y
-
-	def predict_exclusive(self, X):
-		if not np.array_equal(self.X_test, X):
-			self.encoded_X_test = self.clause_bank.prepare_X(X)
-			self.X_test = X.copy()
-
-		Y = np.ascontiguousarray(np.zeros(X.shape[0], dtype=np.uint32))
-
-		for e in range(X.shape[0]):
-			max_class_sum = -self.T
-			max_class = 0
-			clause_outputs = self.clause_bank.calculate_clause_outputs_predict(self.encoded_X_test, e)			
-			for i in range(self.number_of_classes):
-				class_sum = np.dot(self.weight_banks[i].get_weights(), clause_outputs).astype(np.int32)
-				class_sum = np.clip(class_sum, -self.T, self.T)
-				if class_sum > max_class_sum:
-					max_class_sum = class_sum
-					max_class = i
-			Y[e] = max_class
-		return Y
-
-	def clause_precision(self, the_class, positive_polarity, X, Y):
-		clause_outputs = self.transform(X)
-		weights = self.weight_banks[the_class].get_weights()
-		if positive_polarity == 0:
-			positive_clause_outputs = (weights >= 0)[:,np.newaxis].transpose() * clause_outputs
-			true_positive_clause_outputs = clause_outputs[Y==the_class].sum(axis=0)
-			false_positive_clause_outputs = clause_outputs[Y!=the_class].sum(axis=0)
-		else:
-			positive_clause_outputs = (weights < 0)[:,np.newaxis].transpose() * clause_outputs
-			true_positive_clause_outputs = clause_outputs[Y!=the_class].sum(axis=0)
-			false_positive_clause_outputs = clause_outputs[Y==the_class].sum(axis=0)
-		
-		return np.where(true_positive_clause_outputs + false_positive_clause_outputs == 0, 0, 1.0*true_positive_clause_outputs/(true_positive_clause_outputs + false_positive_clause_outputs))
-
-	def clause_recall(self, the_class, positive_polarity, X, Y):
-		clause_outputs = self.transform(X)
-		weights = self.weight_banks[the_class].get_weights()
-		
-		if positive_polarity == 0:
-			positive_clause_outputs = (weights >= 0)[:,np.newaxis].transpose() * clause_outputs
-			true_positive_clause_outputs = positive_clause_outputs[Y==the_class].sum(axis=0)
-		else:
-			positive_clause_outputs = (weights < 0)[:,np.newaxis].transpose() * clause_outputs
-			true_positive_clause_outputs = positive_clause_outputs[Y!=the_class].sum(axis=0)
-			
-		return true_positive_clause_outputs / Y[Y==the_class].shape[0]
-
-	def get_weight(self, the_class, clause):
-		return self.weight_banks[the_class].get_weights()[clause]
-
-	def get_weights(self, the_class):
-		return self.weight_banks[the_class].get_weights()
-
-	def set_weight(self, the_class, clause, weight):
-		self.weight_banks[the_class].get_weights()[clause] = weight
-
-
-class TMMultiChannelClassifier(TMBasis):
-	def __init__(self, number_of_clauses, global_T, T, s, type_i_ii_ratio = 1.0, platform = 'CPU', patch_dim=None, feature_negation=True, boost_true_positive_feedback=1, max_included_literals=None, number_of_state_bits_ta=8, weighted_clauses=False, clause_drop_p = 0.0, literal_drop_p = 0.0):
-		super().__init__(number_of_clauses, T, s, type_i_ii_ratio = type_i_ii_ratio, platform = platform, patch_dim=patch_dim, feature_negation=feature_negation, boost_true_positive_feedback=boost_true_positive_feedback, max_included_literals=max_included_literals, number_of_state_bits_ta=number_of_state_bits_ta, weighted_clauses=weighted_clauses, clause_drop_p = clause_drop_p, literal_drop_p = literal_drop_p)
-		self.global_T = global_T
-
-	def initialize(self, X, Y):
-		self.number_of_classes = int(np.max(Y) + 1)
-	
-		if self.platform == 'CPU':
-			self.clause_bank = ClauseBank(X[0], self.number_of_clauses, self.number_of_state_bits_ta, self.number_of_state_bits_ind, self.patch_dim)
-		elif self.platform == 'CUDA':
-			from tmu.clause_bank_cuda import ClauseBankCUDA
-			self.clause_bank = ClauseBankCUDA(X[0], self.number_of_clauses, self.number_of_state_bits_ta, self.patch_dim)
-		else:
-			print("Unknown Platform")
-			sys.exit(-1)
-
-		self.weight_banks = []
-		for i in range(self.number_of_classes):
-			self.weight_banks.append(WeightBank(np.random.choice([-1,1], size=self.number_of_clauses).astype(np.int32)))
-		
-		self.X_train = {}
-		self.X_test = {}
-		for c in range(X.shape[0]):
-			self.X_train[c] = np.zeros(0, dtype=np.uint32)
-			self.X_test[c] = np.zeros(0, dtype=np.uint32)
-
-		self.encoded_X_train = {}
-		self.encoded_X_test = {}
-
-		if self.max_included_literals == None:
-			self.max_included_literals = self.clause_bank.number_of_literals
-
-	def fit(self, X, Y, shuffle=True):
-		if self.initialized == False:
-			self.initialize(X, Y)
-			self.initialized = True
-
-		for c in range(X.shape[0]):
-			if not np.array_equal(self.X_train[c], X[c]):
-				self.encoded_X_train[c] = self.clause_bank.prepare_X(X[c])
-				self.X_train[c] = X[c].copy()
-
-		Ym = np.ascontiguousarray(Y).astype(np.uint32)
-
-		# Clauses are dropped based on their weights
-		clause_active = np.ones(self.number_of_clauses, dtype=np.uint32)
-		clause_score = np.zeros(self.number_of_clauses, dtype=np.int32)
-		for i in range(self.number_of_classes):
-			clause_score += np.abs(self.weight_banks[i].get_weights())
-		deactivate = np.random.choice(np.arange(self.number_of_clauses), size=int(self.number_of_clauses*self.clause_drop_p), p = clause_score / clause_score.sum())
-		for d in range(deactivate.shape[0]):
-			clause_active[deactivate[d]] = 0
-
-		# Literals are dropped based on their frequency
-		literal_active = (np.zeros(self.clause_bank.number_of_ta_chunks, dtype=np.uint32) | ~0).astype(np.uint32)
-		literal_clause_frequency = self.literal_clause_frequency()
-		if literal_clause_frequency.sum() > 0:
-			deactivate = np.random.choice(np.arange(self.clause_bank.number_of_literals), size=int(self.clause_bank.number_of_literals*self.literal_drop_p), p = literal_clause_frequency / literal_clause_frequency.sum())
-		else:
-			deactivate = np.random.choice(np.arange(self.clause_bank.number_of_literals), size=int(self.clause_bank.number_of_literals*self.literal_drop_p))
-		for d in range(deactivate.shape[0]):
-			ta_chunk = deactivate[d] // 32
-			chunk_pos = deactivate[d] % 32
-			literal_active[ta_chunk] &= (~(1 << chunk_pos))
-
-		if not self.feature_negation:
-			for k in range(self.clause_bank.number_of_literals//2, self.clause_bank.number_of_literals):
-				ta_chunk = k // 32
-				chunk_pos = k % 32
-				literal_active[ta_chunk] &= (~(1 << chunk_pos))
-
-		literal_active = literal_active.astype(np.uint32)
-
-		local_class_sum = np.empty(X.shape[0], dtype=np.int32)
-
-		shuffled_index = np.arange(X.shape[1])
-		if shuffle:
-			np.random.shuffle(shuffled_index)
-
-		for e in shuffled_index:
-			target = Ym[e]
-
-			clause_outputs = []
-			for c in range(X.shape[0]):
-				clause_outputs.append(self.clause_bank.calculate_clause_outputs_update(literal_active, self.encoded_X_train[c], e).copy())
-			
-			global_class_sum = 0
-			for c in range(X.shape[0]):
-				local_class_sum[c] = np.dot(clause_active * self.weight_banks[target].get_weights(), clause_outputs[c]).astype(np.int32)
-				local_class_sum[c] = np.clip(local_class_sum[c], -self.T, self.T)
-				global_class_sum += local_class_sum[c]
-			global_class_sum = np.clip(global_class_sum, -self.global_T[target][0], self.global_T[target][1])
-			global_update_p = 1.0*(self.global_T[target][1] - global_class_sum)/(self.global_T[target][0]+self.global_T[target][1])
-
-			for c in range(X.shape[0]):
-				local_update_p = 1.0*(self.T - local_class_sum[c])/(2*self.T)
-				update_p = np.minimum(local_update_p, global_update_p)
-				self.clause_bank.type_i_feedback(update_p*self.type_i_p, self.s[target], self.boost_true_positive_feedback, clause_active*(self.weight_banks[target].get_weights() >= 0), literal_active, self.encoded_X_train[c], e)
-				self.clause_bank.type_ii_feedback(update_p*self.type_ii_p, clause_active*(self.weight_banks[target].get_weights() < 0), literal_active, self.encoded_X_train[c], e)
-				self.weight_banks[target].increment(clause_outputs[c], update_p, clause_active, True)
-
-			not_target = np.random.randint(self.number_of_classes)
-			while not_target == target:
-				not_target = np.random.randint(self.number_of_classes)
-
-			global_class_sum = 0.0
-			for c in range(X.shape[0]):				
-				local_class_sum[c] = np.dot(clause_active * self.weight_banks[not_target].get_weights(), clause_outputs[c]).astype(np.int32)
-				local_class_sum[c] = np.clip(local_class_sum[c], -self.T, self.T)
-				global_class_sum += local_class_sum[c]
-			global_class_sum = np.clip(global_class_sum, -self.global_T[not_target][0], self.global_T[not_target][1])
-			global_update_p = 1.0*(self.global_T[not_target][0] + global_class_sum)/(self.global_T[not_target][0]+self.global_T[not_target][1])
-
-			for c in range(X.shape[0]):
-				local_update_p = 1.0*(self.T + local_class_sum[c])/(2*self.T)
-				update_p = np.minimum(local_update_p, global_update_p)
-				self.clause_bank.type_i_feedback(update_p*self.type_i_p, self.s[not_target], self.boost_true_positive_feedback, self.max_included_literals, clause_active * (self.weight_banks[not_target].get_weights() < 0), literal_active, self.encoded_X_train[c], e)
-				self.clause_bank.type_ii_feedback(update_p*self.type_ii_p, clause_active*(self.weight_banks[not_target].get_weights() >= 0), literal_active, self.encoded_X_train[c], e)
-				self.weight_banks[not_target].decrement(clause_outputs[c], update_p, clause_active, True)
-		return
-
-	def predict(self, X):
-		for c in range(X.shape[0]):
-			if not np.array_equal(self.X_test[c], X[c]):
-				self.encoded_X_test[c] = self.clause_bank.prepare_X(X[c])
-				self.X_test[c] = X[c].copy()
-
-		Y = np.ascontiguousarray(np.zeros(X.shape[1], dtype=np.uint32))
-
-		for e in range(X.shape[1]):
-			max_class_sum = -maxsize
-			max_class = 0
-
-			clause_outputs = []
-			for c in range(X.shape[0]):
-				clause_outputs.append(self.clause_bank.calculate_clause_outputs_predict(self.encoded_X_test[c], e).copy())
-			
-			for i in range(self.number_of_classes):
-				global_class_sum = 1
-				for c in range(X.shape[0]):
-					local_class_sum = np.dot(self.weight_banks[i].get_weights(), clause_outputs[c]).astype(np.int32)
-					local_class_sum = np.clip(local_class_sum, -self.T, self.T)
-					global_class_sum *= local_class_sum >= 0
-				global_class_sum = np.clip(global_class_sum, -self.global_T[i][0], self.global_T[i][1])
-
-				if global_class_sum > max_class_sum:
-					max_class_sum = global_class_sum
-					max_class = i
-			Y[e] = max_class
-		
-		return Y
-
-	def clause_precision(self, the_class, positive_polarity, X, Y):
-		clause_outputs = self.transform(X)
-		weights = self.weight_banks[the_class].get_weights()
-		if positive_polarity == 0:
-			positive_clause_outputs = (weights >= 0)[:,np.newaxis].transpose() * clause_outputs
-			true_positive_clause_outputs = clause_outputs[Y==the_class].sum(axis=0)
-			false_positive_clause_outputs = clause_outputs[Y!=the_class].sum(axis=0)
-		else:
-			positive_clause_outputs = (weights < 0)[:,np.newaxis].transpose() * clause_outputs
-			true_positive_clause_outputs = clause_outputs[Y!=the_class].sum(axis=0)
-			false_positive_clause_outputs = clause_outputs[Y==the_class].sum(axis=0)
-		
-		return np.where(true_positive_clause_outputs + false_positive_clause_outputs == 0, 0, 1.0*true_positive_clause_outputs/(true_positive_clause_outputs + false_positive_clause_outputs))
-
-	def clause_recall(self, the_class, positive_polarity, X, Y):
-		clause_outputs = self.transform(X)
-		weights = self.weight_banks[the_class].get_weights()
-		
-		if positive_polarity == 0:
-			positive_clause_outputs = (weights >= 0)[:,np.newaxis].transpose() * clause_outputs
-			true_positive_clause_outputs = positive_clause_outputs[Y==the_class].sum(axis=0)
-		else:
-			positive_clause_outputs = (weights < 0)[:,np.newaxis].transpose() * clause_outputs
-			true_positive_clause_outputs = positive_clause_outputs[Y!=the_class].sum(axis=0)
-			
-		return true_positive_clause_outputs / Y[Y==the_class].shape[0]
-
-	def get_weight(self, the_class, clause):
-		return self.weight_banks[the_class].get_weights()[clause]
-
-	def set_weight(self, the_class, clause, weight):
-		self.weight_banks[the_class].get_weights()[clause] = weight
-
-class TMOneVsOneClassifier(TMBasis):
-	def __init__(self, number_of_clauses, T, s, type_i_ii_ratio = 1.0, platform = 'CPU', patch_dim=None, feature_negation=True, boost_true_positive_feedback=1, max_included_literals=None, number_of_state_bits_ta=8, weighted_clauses=False, clause_drop_p = 0.0, literal_drop_p = 0.0):
-		super().__init__(number_of_clauses, T, s, type_i_ii_ratio = type_i_ii_ratio, platform = platform, patch_dim=patch_dim, feature_negation=feature_negation, boost_true_positive_feedback=boost_true_positive_feedback, max_included_literals=max_included_literals, number_of_state_bits_ta=number_of_state_bits_ta, weighted_clauses=weighted_clauses, clause_drop_p = clause_drop_p, literal_drop_p = literal_drop_p)
-
-	def initialize(self, X, Y):
-		self.number_of_classes = int(np.max(Y) + 1)
-		self.number_of_outputs = self.number_of_classes * (self.number_of_classes-1)
-
-		if self.platform == 'CPU':
-			self.clause_bank = ClauseBank(X, self.number_of_clauses, self.number_of_state_bits_ta, self.number_of_state_bits_ind, self.patch_dim)
-		elif self.platform == 'CUDA':
-			from tmu.clause_bank_cuda import ClauseBankCUDA
-			self.clause_bank = ClauseBankCUDA(X, self.number_of_clauses, self.number_of_state_bits_ta, self.patch_dim)
-		else:
-			print("Unknown Platform")
-			sys.exit(-1)
-
-		self.weight_banks = []
-		for i in range(self.number_of_outputs):
-			self.weight_banks.append(WeightBank(np.ones(self.number_of_clauses).astype(np.int32)))
-		
-		if self.max_included_literals == None:
-			self.max_included_literals = self.clause_bank.number_of_literals
-
-	def fit(self, X, Y, shuffle=True):
-		if self.initialized == False:
-			self.initialize(X, Y)
-			self.initialized = True
-
-		if not np.array_equal(self.X_train, X):
-			self.encoded_X_train = self.clause_bank.prepare_X(X)
-			self.X_train = X.copy()
-
-		Ym = np.ascontiguousarray(Y).astype(np.uint32)
-		
-		clause_active = np.ascontiguousarray(np.random.choice(2, self.number_of_clauses, p=[self.clause_drop_p, 1.0 - self.clause_drop_p]).astype(np.int32))
-		literal_active = (np.zeros(self.clause_bank.number_of_ta_chunks, dtype=np.uint32) | ~0).astype(np.uint32)
-		if not self.feature_negation:
-			for k in range(self.clause_bank.number_of_literals//2, self.clause_bank.number_of_literals):
-				ta_chunk = k // 32
-				chunk_pos = k % 32
-				literal_active[ta_chunk] &= (~(1 << chunk_pos))
-
-		literal_active = literal_active.astype(np.uint32)
-
-		shuffled_index = np.arange(X.shape[0])
-		if shuffle:
-			np.random.shuffle(shuffled_index)
-
-		for e in shuffled_index:
-			clause_outputs = self.clause_bank.calculate_clause_outputs_update(literal_active, self.encoded_X_train, e)
-			
-			target = Ym[e]
-			not_target = np.random.randint(self.number_of_classes)
-			while not_target == target:
-				not_target = np.random.randint(self.number_of_classes)
-
-			output = target * (self.number_of_classes-1) + not_target - (not_target > target)
-
-			class_sum = np.dot(clause_active * self.weight_banks[output].get_weights(), clause_outputs).astype(np.int32)
-			class_sum = np.clip(class_sum, -self.T, self.T)
-			update_p = (self.T - class_sum)/(2*self.T)
-
-			self.clause_bank.type_i_feedback(update_p*self.type_i_p, self.s, self.boost_true_positive_feedback, self.max_included_literals, clause_active*(self.weight_banks[output].get_weights() >= 0), literal_active, self.encoded_X_train, e)
-			self.clause_bank.type_ii_feedback(update_p*self.type_ii_p, clause_active*(self.weight_banks[output].get_weights() < 0), literal_active, self.encoded_X_train, e)
-			self.weight_banks[output].increment(clause_outputs, update_p, clause_active, True)
-
-			output = not_target * (self.number_of_classes-1) + target - (target > not_target)
-
-			class_sum = np.dot(clause_active * self.weight_banks[output].get_weights(), clause_outputs).astype(np.int32)
-			class_sum = np.clip(class_sum, -self.T, self.T)
-			update_p = (self.T + class_sum)/(2*self.T)
-		
-			self.clause_bank.type_i_feedback(update_p*self.type_i_p, self.s, self.boost_true_positive_feedback, self.max_included_literals, clause_active * (self.weight_banks[output].get_weights() < 0), literal_active, self.encoded_X_train, e)
-			self.clause_bank.type_ii_feedback(update_p*self.type_ii_p, clause_active*(self.weight_banks[output].get_weights() >= 0), literal_active, self.encoded_X_train, e)
-			self.weight_banks[output].decrement(clause_outputs, update_p, clause_active, True)
-		return
-
-	def predict(self, X):
-		if not np.array_equal(self.X_test, X):
-			self.encoded_X_test = self.clause_bank.prepare_X(X)
-			self.X_test = X.copy()
-
-		Y = np.ascontiguousarray(np.zeros(X.shape[0], dtype=np.uint32))
-
-		for e in range(X.shape[0]):
-			clause_outputs = self.clause_bank.calculate_clause_outputs_predict(self.encoded_X_test, e)
-
-			max_class_sum = -self.T*self.number_of_classes
-			max_class = 0
-			for i in range(self.number_of_classes):
-				class_sum = 0
-				for output in range(i * (self.number_of_classes - 1), (i+1) * (self.number_of_classes-1)):
-					output_sum = np.dot(self.weight_banks[output].get_weights(), clause_outputs).astype(np.int32)
-					output_sum = np.clip(output_sum, -self.T, self.T)
-					class_sum += output_sum
-
-				if class_sum > max_class_sum:
-					max_class_sum = class_sum
-					max_class = i
-			Y[e] = max_class
-		return Y
-	
-	def clause_precision(self, the_class, positive_polarity, X, Y):
-		clause_outputs = self.transform(X)
-		precision = np.zeros((self.number_of_classes - 1, self.number_of_clauses))
-		for i in range(self.number_of_classes - 1):
-			other_class = i + (i >= the_class)
-			output = the_class * (self.number_of_classes - 1) + i
-			weights = self.weight_banks[output].get_weights()
-			if positive_polarity:
-				positive_clause_outputs = (weights >= 0)[:,np.newaxis].transpose() * clause_outputs
-				true_positive_clause_outputs = positive_clause_outputs[Y==the_class].sum(axis=0)
-				false_positive_clause_outputs = positive_clause_outputs[Y==other_class].sum(axis=0)
-				precision[i] = np.where(true_positive_clause_outputs + false_positive_clause_outputs == 0, 0, true_positive_clause_outputs/(true_positive_clause_outputs + false_positive_clause_outputs))
-			else:
-				positive_clause_outputs = (weights < 0)[:,np.newaxis].transpose() * clause_outputs
-				true_positive_clause_outputs = positive_clause_outputs[Y==other_class].sum(axis=0)
-				false_positive_clause_outputs = positive_clause_outputs[Y==the_class].sum(axis=0)
-				precision[i] = np.where(true_positive_clause_outputs + false_positive_clause_outputs == 0, 0, true_positive_clause_outputs/(true_positive_clause_outputs + false_positive_clause_outputs))
-	
-		return precision
-
-	def clause_recall(self, the_class, positive_polarity, X, Y):
-		clause_outputs = self.transform(X)
-		recall = np.zeros((self.number_of_classes - 1, self.number_of_clauses))
-		for i in range(self.number_of_classes - 1):
-			other_class = i + (i >= the_class)
-			output = the_class * (self.number_of_classes - 1) + i
-			weights = self.weight_banks[output].get_weights()
-			if positive_polarity:
-				positive_clause_outputs = (weights >= 0)[:,np.newaxis].transpose() * clause_outputs
-				true_positive_clause_outputs = positive_clause_outputs[Y==the_class].sum(axis=0)
-				recall[i] =	true_positive_clause_outputs / Y[Y==the_class].shape[0]
-			else:
-				positive_clause_outputs = (weights < 0)[:,np.newaxis].transpose() * clause_outputs
-				true_positive_clause_outputs = positive_clause_outputs[Y==other_class].sum(axis=0)
-				recall[i] = true_positive_clause_outputs / Y[Y==other_class].shape[0]
-		return recall
-
-	def get_weight(self, output, clause):
-		return self.weight_banks[output].get_weights()[clause]
-
-	def set_weight(self, output, weight):
-		self.weight_banks[output].get_weights()[output] = weight
-
-class TMRegressor(TMBasis):
-	def __init__(self, number_of_clauses, T, s, platform='CPU', patch_dim=None, feature_negation=True, boost_true_positive_feedback=1, max_included_literals=None, number_of_state_bits_ta=8, weighted_clauses=False, clause_drop_p = 0.0, literal_drop_p = 0.0):
-		super().__init__(number_of_clauses, T, s, platform=platform, patch_dim=patch_dim, feature_negation=feature_negation, boost_true_positive_feedback=boost_true_positive_feedback, max_included_literals=max_included_literals, number_of_state_bits_ta=number_of_state_bits_ta, weighted_clauses=weighted_clauses, clause_drop_p = clause_drop_p, literal_drop_p = literal_drop_p)
-
-	def initialize(self, X, Y):
-		self.max_y = np.max(Y)
-		self.min_y = np.min(Y)
-
-		if self.platform == 'CPU':
-			self.clause_bank = ClauseBank(X, self.number_of_clauses, self.number_of_state_bits_ta, self.number_of_state_bits_ind, self.patch_dim)
-		elif self.platform == 'CUDA':
-			from tmu.clause_bank_cuda import ClauseBankCUDA
-			self.clause_bank = ClauseBankCUDA(X, self.number_of_clauses, self.number_of_state_bits_ta, self.patch_dim)
-		else:
-			print("Unknown Platform")
-			sys.exit(-1)
-			
-		self.weight_bank = WeightBank(np.ones(self.number_of_clauses).astype(np.int32))
-
-		if self.max_included_literals == None:
-			self.max_included_literals = self.clause_bank.number_of_literals
-
-	def fit(self, X, Y, shuffle=True):
-		if self.initialized == False:
-			self.initialize(X, Y)
-			self.initialized = True
-
-		if not np.array_equal(self.X_train, X):
-			self.encoded_X_train = self.clause_bank.prepare_X(X)
-			self.X_train = X.copy()
-		encoded_Y = np.ascontiguousarray(((Y - self.min_y)/(self.max_y - self.min_y)*self.T).astype(np.int32))
-
-		clause_active = np.ascontiguousarray(np.random.choice(2, self.number_of_clauses, p=[self.clause_drop_p, 1.0 - self.clause_drop_p]).astype(np.int32))
-		literal_active = (np.zeros(self.clause_bank.number_of_ta_chunks, dtype=np.uint32) | ~0).astype(np.uint32)
-		if not self.feature_negation:
-			for k in range(self.clause_bank.number_of_literals//2, self.clause_bank.number_of_literals):
-				ta_chunk = k // 32
-				chunk_pos = k % 32
-				literal_active[ta_chunk] &= (~(1 << chunk_pos))
-
-		literal_active = literal_active.astype(np.uint32)
-
-		shuffled_index = np.arange(X.shape[0])
-		if shuffle:
-			np.random.shuffle(shuffled_index)
-
-		for e in shuffled_index:
-			clause_outputs = self.clause_bank.calculate_clause_outputs_update(literal_active, self.encoded_X_train, e)
-
-			pred_y = np.dot(clause_active * self.weight_bank.get_weights(), clause_outputs).astype(np.int32)
-			pred_y = np.clip(pred_y, 0, self.T)
-			prediction_error = pred_y - encoded_Y[e]; 
-
-			update_p = (1.0*prediction_error/self.T)**2
-
-			if pred_y < encoded_Y[e]:
-				self.clause_bank.type_i_feedback(update_p, self.s, self.boost_true_positive_feedback, self.max_included_literals, clause_active, literal_active, self.encoded_X_train, e)
-				if self.weighted_clauses:
-					self.weight_bank.increment(clause_outputs, update_p, clause_active, False)
-			elif pred_y > encoded_Y[e]:
-				self.clause_bank.type_ii_feedback(update_p, clause_active, literal_active, self.encoded_X_train, e)
-				if self.weighted_clauses:
-					self.weight_bank.decrement(clause_outputs, update_p, clause_active, False)
-		return
-
-	def predict(self, X):
-		if not np.array_equal(self.X_test, X):
-			self.encoded_X_test = self.clause_bank.prepare_X(X)
-			self.X_test = X.copy()
-
-		Y = np.ascontiguousarray(np.zeros(X.shape[0]))
-		for e in range(X.shape[0]):
-			clause_outputs = self.clause_bank.calculate_clause_outputs_predict(self.encoded_X_test, e)
-			pred_y = np.dot(self.weight_bank.get_weights(), clause_outputs).astype(np.int32)
-			Y[e] = 1.0*pred_y * (self.max_y - self.min_y)/(self.T) + self.min_y
-		return Y
-
-	def get_weight(self, clause):
-		return self.weight_bank.get_weights()[clause]
-
-	def set_weight(self, clause, weight):
-		self.weight_banks.get_weights()[clause] = weight
-
diff --git a/tmu/weight_bank.py b/tmu/weight_bank.py
index 743280ff..2563bd97 100644
--- a/tmu/weight_bank.py
+++ b/tmu/weight_bank.py
@@ -21,26 +21,27 @@
 # This code implements the Convolutional Tsetlin Machine from paper arXiv:1905.09688
 # https://arxiv.org/abs/1905.09688
 
-from ._wb import ffi, lib
+from tmu.tmulib import ffi, lib
 
 import numpy as np
 
-class WeightBank():
-	def __init__(self, weights):
-		self.number_of_clauses = weights.shape[0]
 
-		self.weights = np.ascontiguousarray(weights, dtype=np.int32)
-		self.cw_p = ffi.cast("int *", self.weights.ctypes.data)
+class WeightBank:
+    def __init__(self, weights):
+        self.number_of_clauses = weights.shape[0]
 
-	def increment(self, clause_output, update_p, clause_active, positive_weights):
-		co_p = ffi.cast("unsigned int *", clause_output.ctypes.data)
-		ca_p = ffi.cast("unsigned int *", clause_active.ctypes.data)
-		lib.wb_increment(self.cw_p, self.number_of_clauses, co_p, update_p, ca_p, int(positive_weights))
+        self.weights = np.ascontiguousarray(weights, dtype=np.int32)
+        self.cw_p = ffi.cast("int *", self.weights.ctypes.data)
 
-	def decrement(self, clause_output, update_p, clause_active, negative_weights):
-		co_p = ffi.cast("unsigned int *", clause_output.ctypes.data)
-		ca_p = ffi.cast("unsigned int *", clause_active.ctypes.data)
-		lib.wb_decrement(self.cw_p, self.number_of_clauses, co_p, update_p, ca_p, int(negative_weights))
+    def increment(self, clause_output, update_p, clause_active, positive_weights):
+        co_p = ffi.cast("unsigned int *", clause_output.ctypes.data)
+        ca_p = ffi.cast("unsigned int *", clause_active.ctypes.data)
+        lib.wb_increment(self.cw_p, self.number_of_clauses, co_p, update_p, ca_p, int(positive_weights))
 
-	def get_weights(self):
-		return self.weights
+    def decrement(self, clause_output, update_p, clause_active, negative_weights):
+        co_p = ffi.cast("unsigned int *", clause_output.ctypes.data)
+        ca_p = ffi.cast("unsigned int *", clause_active.ctypes.data)
+        lib.wb_decrement(self.cw_p, self.number_of_clauses, co_p, update_p, ca_p, int(negative_weights))
+
+    def get_weights(self):
+        return self.weights
diff --git a/tmu/weight_bank_extension_build.py b/tmu/weight_bank_extension_build.py
deleted file mode 100644
index b6eea826..00000000
--- a/tmu/weight_bank_extension_build.py
+++ /dev/null
@@ -1,42 +0,0 @@
-# Copyright (c) 2023 Ole-Christoffer Granmo
-
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-
-# The above copyright notice and this permission notice shall be included in all
-# copies or substantial portions of the Software.
-
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-# SOFTWARE.
-
-# This code implements the Convolutional Tsetlin Machine from paper arXiv:1905.09688
-# https://arxiv.org/abs/1905.09688
-
-from cffi import FFI
-import pathlib
-current_dir = pathlib.Path(__file__).parent
-
-ffibuilder = FFI()
-
-with current_dir.joinpath("WeightBank.h").open("r") as f:
-    ffibuilder.cdef(f.read())
-
-with current_dir.joinpath("WeightBank.c").open("r") as f:
-    ffibuilder.set_source(
-        "tmu._wb",
-        f.read(),
-        include_dirs=[current_dir.absolute()],
-        source_extension=".c",
-    )
-
-if __name__ == "__main__":
-    ffibuilder.compile(verbose=True)
\ No newline at end of file


Test patch:
diff --git a/.github/workflows/build_test.yaml b/.github/workflows/build_test.yaml
new file mode 100644
index 00000000..669d36cb
--- /dev/null
+++ b/.github/workflows/build_test.yaml
@@ -0,0 +1,40 @@
+name: Python package
+
+on: [push]
+
+jobs:
+  build:
+
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-20.04, ubuntu-22.04, macos-latest, windows-latest]
+        python-version: ["3.6", "3.7", "3.8", "3.9", "3.10"]
+        exclude:
+          - os: ubuntu-22.04
+            python-version: "3.6"
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - 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 -e .
+      - name: Test with pytest
+        run: pytest test --doctest-modules --junitxml=junit/test-results-${{ matrix.os }}-${{ matrix.python-version }}.xml
+      - name: Upload pytest test results
+        uses: actions/upload-artifact@v3
+        with:
+          name: pytest-results-${{ matrix.os }}-${{ matrix.python-version }}
+          path: junit/test-results-${{ matrix.os }}-${{ matrix.python-version }}.xml
+        # Use always() to always run this step to publish test results when there are test failures
+        if: ${{ always() }}
diff --git a/test/__init__.py b/test/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/test/test_datasets.py b/test/test_datasets.py
new file mode 100644
index 00000000..4116f4a6
--- /dev/null
+++ b/test/test_datasets.py
@@ -0,0 +1,27 @@
+import pytest
+class TestMNISTDatasetobject:
+
+    def setup_class(self):
+        from tmu.datasets import MNIST
+        self.dataset_instance = MNIST()
+
+    def teardown_class(self):
+        self.dataset_instance = None
+
+    def test_mnist_dict(self):
+        dataset = self.dataset_instance.retrieve_dataset()
+        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
+
+    def test_mnist_list(self):
+        dataset = self.dataset_instance.get_list()
+        assert len(dataset) == 4
+        assert isinstance(dataset, list)
+
diff --git a/test/test_imports.py b/test/test_imports.py
new file mode 100644
index 00000000..beb73f1f
--- /dev/null
+++ b/test/test_imports.py
@@ -0,0 +1,5 @@
+
+
+
+def test_import_tmulib():
+    import tmu.tmulib
\ No newline at end of file