diff --git a/tests/factors/test_covariances.py b/tests/factors/test_covariances.py
index 114fe0d..f1eba04 100644
--- a/tests/factors/test_covariances.py
+++ b/tests/factors/test_covariances.py
@@ -8,7 +8,7 @@
 
 from kronfluence.analyzer import Analyzer, prepare_model
 from kronfluence.arguments import FactorArguments
-from kronfluence.module.constants import (
+from kronfluence.utils.constants import (
     ACTIVATION_COVARIANCE_MATRIX_NAME,
     COVARIANCE_FACTOR_NAMES,
     GRADIENT_COVARIANCE_MATRIX_NAME,
diff --git a/tests/factors/test_eigens.py b/tests/factors/test_eigens.py
index df19e64..39544e6 100644
--- a/tests/factors/test_eigens.py
+++ b/tests/factors/test_eigens.py
@@ -8,7 +8,7 @@
 
 from kronfluence.analyzer import Analyzer, prepare_model
 from kronfluence.arguments import FactorArguments
-from kronfluence.module.constants import (
+from kronfluence.utils.constants import (
     ACTIVATION_EIGENVECTORS_NAME,
     EIGENDECOMPOSITION_FACTOR_NAMES,
     GRADIENT_EIGENVECTORS_NAME,
diff --git a/tests/gpu_tests/compile_test.py b/tests/gpu_tests/compile_test.py
index cf3bc4e..61e9db5 100644
--- a/tests/gpu_tests/compile_test.py
+++ b/tests/gpu_tests/compile_test.py
@@ -8,7 +8,7 @@
 
 from kronfluence.analyzer import Analyzer, prepare_model
 from kronfluence.arguments import FactorArguments, ScoreArguments
-from kronfluence.module.constants import (
+from kronfluence.utils.constants import (
     ALL_MODULE_NAME,
     COVARIANCE_FACTOR_NAMES,
     LAMBDA_FACTOR_NAMES,
diff --git a/tests/gpu_tests/cpu_test.py b/tests/gpu_tests/cpu_test.py
index 535ac9e..455327d 100644
--- a/tests/gpu_tests/cpu_test.py
+++ b/tests/gpu_tests/cpu_test.py
@@ -8,7 +8,7 @@
 
 from kronfluence.analyzer import Analyzer, prepare_model
 from kronfluence.arguments import FactorArguments, ScoreArguments
-from kronfluence.module.constants import (
+from kronfluence.utils.constants import (
     ALL_MODULE_NAME,
     COVARIANCE_FACTOR_NAMES,
     LAMBDA_FACTOR_NAMES,
diff --git a/tests/gpu_tests/ddp_test.py b/tests/gpu_tests/ddp_test.py
index 37711cd..3cad50d 100644
--- a/tests/gpu_tests/ddp_test.py
+++ b/tests/gpu_tests/ddp_test.py
@@ -11,7 +11,7 @@
 
 from kronfluence.analyzer import Analyzer, prepare_model
 from kronfluence.arguments import FactorArguments, ScoreArguments
-from kronfluence.module.constants import (
+from kronfluence.utils.constants import (
     ALL_MODULE_NAME,
     COVARIANCE_FACTOR_NAMES,
     LAMBDA_FACTOR_NAMES,
diff --git a/tests/gpu_tests/fsdp_test.py b/tests/gpu_tests/fsdp_test.py
index 4d95898..9a594d0 100644
--- a/tests/gpu_tests/fsdp_test.py
+++ b/tests/gpu_tests/fsdp_test.py
@@ -15,7 +15,7 @@
 
 from kronfluence.analyzer import Analyzer, prepare_model
 from kronfluence.arguments import FactorArguments, ScoreArguments
-from kronfluence.module.constants import (
+from kronfluence.utils.constants import (
     ALL_MODULE_NAME,
     COVARIANCE_FACTOR_NAMES,
     LAMBDA_FACTOR_NAMES,
diff --git a/tests/gpu_tests/test_offload_cpu.py b/tests/gpu_tests/test_offload_cpu.py
index 521dfe1..5b8522a 100644
--- a/tests/gpu_tests/test_offload_cpu.py
+++ b/tests/gpu_tests/test_offload_cpu.py
@@ -8,7 +8,7 @@
 
 from kronfluence.analyzer import Analyzer, prepare_model
 from kronfluence.arguments import FactorArguments, ScoreArguments
-from kronfluence.module.constants import ALL_MODULE_NAME
+from kronfluence.utils.constants import ALL_MODULE_NAME
 from kronfluence.utils.dataset import DataLoaderKwargs
 from tests.utils import prepare_test
 
diff --git a/tests/scores/test_pairwise_scores.py b/tests/scores/test_pairwise_scores.py
index b506025..3e04016 100644
--- a/tests/scores/test_pairwise_scores.py
+++ b/tests/scores/test_pairwise_scores.py
@@ -8,7 +8,7 @@
 
 from kronfluence.analyzer import Analyzer, prepare_model
 from kronfluence.arguments import FactorArguments, ScoreArguments
-from kronfluence.module.constants import ALL_MODULE_NAME
+from kronfluence.utils.constants import ALL_MODULE_NAME
 from kronfluence.task import Task
 from kronfluence.utils.dataset import DataLoaderKwargs
 from tests.utils import ATOL, RTOL, check_tensor_dict_equivalence, prepare_test
diff --git a/tests/scores/test_self_scores.py b/tests/scores/test_self_scores.py
index cb16382..4803706 100644
--- a/tests/scores/test_self_scores.py
+++ b/tests/scores/test_self_scores.py
@@ -8,7 +8,7 @@
 
 from kronfluence.analyzer import Analyzer, prepare_model
 from kronfluence.arguments import FactorArguments, ScoreArguments
-from kronfluence.module.constants import ALL_MODULE_NAME
+from kronfluence.utils.constants import ALL_MODULE_NAME
 from kronfluence.task import Task
 from kronfluence.utils.dataset import DataLoaderKwargs
 from tests.utils import ATOL, RTOL, check_tensor_dict_equivalence, prepare_test
diff --git a/tests/test_per_sample_gradients.py b/tests/test_per_sample_gradients.py
index a5d333a..18c8cd8 100644
--- a/tests/test_per_sample_gradients.py
+++ b/tests/test_per_sample_gradients.py
@@ -13,7 +13,7 @@
 
 from kronfluence.analyzer import Analyzer, prepare_model
 from kronfluence.arguments import FactorArguments
-from kronfluence.module.constants import (
+from kronfluence.utils.constants import (
     LAMBDA_MATRIX_NAME,
     PRECONDITIONED_GRADIENT_NAME,
 )
diff --git a/tests/testable_tasks/classification.py b/tests/testable_tasks/classification.py
index 819d5eb..a8d18c2 100644
--- a/tests/testable_tasks/classification.py
+++ b/tests/testable_tasks/classification.py
@@ -6,7 +6,7 @@
 import torch.nn.functional as F
 import torchvision
 from accelerate.utils import set_seed
-from task import Task
+from kronfluence.task import Task
 from torch import nn
 from torch.utils import data
 
diff --git a/tests/testable_tasks/language_modeling.py b/tests/testable_tasks/language_modeling.py
index fdc2ac8..39cfd26 100644
--- a/tests/testable_tasks/language_modeling.py
+++ b/tests/testable_tasks/language_modeling.py
@@ -6,7 +6,7 @@
 import torch
 import torch.nn.functional as F
 from datasets import load_dataset
-from task import Task
+from kronfluence.task import Task
 from torch import nn
 from torch.utils import data
 from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, Conv1D
diff --git a/tests/testable_tasks/regression.py b/tests/testable_tasks/regression.py
index b4e5eec..6b9e19b 100644
--- a/tests/testable_tasks/regression.py
+++ b/tests/testable_tasks/regression.py
@@ -1,7 +1,7 @@
 # pylint: skip-file
 
 import math
-from typing import Dict, Tuple
+from typing import Tuple
 
 import torch
 import torch.nn.functional as F
diff --git a/tests/testable_tasks/text_classification.py b/tests/testable_tasks/text_classification.py
index 806c522..893c546 100644
--- a/tests/testable_tasks/text_classification.py
+++ b/tests/testable_tasks/text_classification.py
@@ -5,7 +5,7 @@
 import torch
 import torch.nn.functional as F
 from datasets import load_dataset
-from task import Task
+from kronfluence.task import Task
 from torch import nn
 from torch.utils import data
 from transformers import AutoConfig, AutoModelForSequenceClassification, AutoTokenizer