diff --git a/src/__init__.py b/src/__init__.py
index 206d45f..af42cce 100644
--- a/src/__init__.py
+++ b/src/__init__.py
@@ -2,3 +2,4 @@
 from .quantum_gate import gates
 from .visualize import *
 from .tools import *
+from .qubit import *
diff --git a/src/quantum_circuit/base/base_core.py b/src/quantum_circuit/base/base_core.py
index 58ab3a3..4504a49 100644
--- a/src/quantum_circuit/base/base_core.py
+++ b/src/quantum_circuit/base/base_core.py
@@ -1,5 +1,5 @@
 import numpy as np
-from ...tools import qubit
+from ...qubit import qubit
 from ..interface import CoreInterface
 
 
diff --git a/src/quantum_circuit/gpu/gpu_core.py b/src/quantum_circuit/gpu/gpu_core.py
index 969fd8b..12b50ef 100644
--- a/src/quantum_circuit/gpu/gpu_core.py
+++ b/src/quantum_circuit/gpu/gpu_core.py
@@ -1,6 +1,6 @@
 import cupy as cp
 from ..interface import CoreInterface
-from ...tools import qubit
+from ...qubit import qubit
 
 
 class GpuCore(CoreInterface):
diff --git a/src/quantum_circuit/gpu_sparse/gpu_sparse_core.py b/src/quantum_circuit/gpu_sparse/gpu_sparse_core.py
index 84caa7f..a92ecca 100644
--- a/src/quantum_circuit/gpu_sparse/gpu_sparse_core.py
+++ b/src/quantum_circuit/gpu_sparse/gpu_sparse_core.py
@@ -2,7 +2,7 @@
 from cupyx.scipy.sparse import csr_matrix, csc_matrix
 from cupyx.scipy.sparse import kron
 from ..interface import CoreInterface
-from ...tools import qubit
+from ...qubit import qubit
 
 
 class GpuSparseCore(CoreInterface):
diff --git a/src/quantum_circuit/sparse/sparse_core.py b/src/quantum_circuit/sparse/sparse_core.py
index 602bd3f..dd5f383 100644
--- a/src/quantum_circuit/sparse/sparse_core.py
+++ b/src/quantum_circuit/sparse/sparse_core.py
@@ -1,6 +1,6 @@
 from scipy import sparse as sp
 from ..interface import CoreInterface
-from ...tools import qubit
+from ...qubit import qubit
 
 
 class SparseCore(CoreInterface):
diff --git a/src/qubit/__init__.py b/src/qubit/__init__.py
new file mode 100644
index 0000000..3b012dd
--- /dev/null
+++ b/src/qubit/__init__.py
@@ -0,0 +1 @@
+from .qubit import qubit
diff --git a/src/tools/qubit.py b/src/qubit/qubit.py
similarity index 100%
rename from src/tools/qubit.py
rename to src/qubit/qubit.py
diff --git a/src/tools/__init__.py b/src/tools/__init__.py
index ccc6fcd..99e81a8 100644
--- a/src/tools/__init__.py
+++ b/src/tools/__init__.py
@@ -1,4 +1,3 @@
-from .qubit import qubit
 from .measure import measure
 from .probability import probability
 from .phase_angle import phase_angle as phaseangle
diff --git a/src/tools/probability.py b/src/tools/probability.py
index 8e6dbe7..5df8d5f 100644
--- a/src/tools/probability.py
+++ b/src/tools/probability.py
@@ -1,20 +1,25 @@
-from numpy import abs, around, multiply, square, log2
+from numpy import abs, around, multiply, square, log2, ndarray
 from .base import convert_state
+from typing import Union
+from ..quantum_circuit import QuantumCircuit
 
 
 def probability(
-    quantumstate, show_percent: bool = False, show_bit: int = -1, round: int = 3
+    quantum_state: Union[ndarray, QuantumCircuit],
+    show_percent: bool = False,
+    show_bit: int = -1,
+    round: int = 3,
 ):
 
-    circuit_size = int(log2(quantumstate.size))
-    quantumstate = convert_state(quantumstate)
+    circuit_size = int(log2(quantum_state.size))
+    quantum_state = convert_state(quantum_state)
     if round < 0:
         exit(
             f"Error: QuantumCircuit.tools.probabilities() -- round placement must be a value greater than 0."
         )
 
     if isinstance(show_bit, int) and show_bit < 0:
-        probability = abs(square(quantumstate))
+        probability = abs(square(quantum_state))
 
     elif isinstance(show_bit, str) or isinstance(show_bit, int):
 
@@ -22,7 +27,7 @@ def probability(
             exit(
                 f"Error: QuantumCircuit.tools.probabilities() -- Called bit to find phase angle is not within range of possible values."
             )
-        probability = abs(square(quantumstate[show_bit]))
+        probability = abs(square(quantum_state[show_bit]))
 
     if show_percent:
         probability = multiply(probability, 100)