diff --git a/examples/LennardJones/LJ_data.py b/examples/LennardJones/LJ_data.py
index 6226ff6f8..e3d9249da 100644
--- a/examples/LennardJones/LJ_data.py
+++ b/examples/LennardJones/LJ_data.py
@@ -32,9 +32,9 @@
 mpi4py.rc.threads = False
 
 # HydraGNN
-from hydragnn.utils.abstractrawdataset import AbstractBaseDataset
-from hydragnn.utils import nsplit
-from hydragnn.preprocess.utils import get_radius_graph_pbc
+from hydragnn.utils.datasets.abstractrawdataset import AbstractBaseDataset
+from hydragnn.utils.distributed import nsplit
+from hydragnn.preprocess.graph_samples_checks_and_updates import get_radius_graph_pbc
 
 # Angstrom unit
 primitive_bravais_lattice_constant_x = 3.8
diff --git a/examples/LennardJones/LJ_inference_plots.py b/examples/LennardJones/LJ_inference_plots.py
index 324da425f..9f5d00cac 100644
--- a/examples/LennardJones/LJ_inference_plots.py
+++ b/examples/LennardJones/LJ_inference_plots.py
@@ -22,11 +22,11 @@
 import numpy as np
 
 import hydragnn
-from hydragnn.utils.time_utils import Timer
+from hydragnn.utils.profiling_and_tracing.time_utils import Timer
 from hydragnn.utils.distributed import get_device
 from hydragnn.utils.model import load_existing_model
-from hydragnn.utils.pickledataset import SimplePickleDataset
-from hydragnn.utils.config_utils import (
+from hydragnn.utils.datasets.pickledataset import SimplePickleDataset
+from hydragnn.utils.input_config_parsing.config_utils import (
     update_config,
 )
 from hydragnn.models.create import create_model_config
@@ -35,7 +35,7 @@
 from scipy.interpolate import griddata
 
 try:
-    from hydragnn.utils.adiosdataset import AdiosWriter, AdiosDataset
+    from hydragnn.utils.datasets.adiosdataset import AdiosWriter, AdiosDataset
 except ImportError:
     pass
 
diff --git a/examples/LennardJones/LennardJones.py b/examples/LennardJones/LennardJones.py
index 045b1d251..ecaa06629 100644
--- a/examples/LennardJones/LennardJones.py
+++ b/examples/LennardJones/LennardJones.py
@@ -30,16 +30,19 @@
 
 # HydraGNN
 import hydragnn
-from hydragnn.utils.print_utils import log
-from hydragnn.utils.time_utils import Timer
-import hydragnn.utils.tracer as tr
+from hydragnn.utils.print.print_utils import log
+from hydragnn.utils.profiling_and_tracing.time_utils import Timer
+import hydragnn.utils.profiling_and_tracing.tracer as tr
 from hydragnn.preprocess.load_data import split_dataset
-from hydragnn.utils.distdataset import DistDataset
-from hydragnn.utils.pickledataset import SimplePickleWriter, SimplePickleDataset
-from hydragnn.preprocess.utils import gather_deg
+from hydragnn.utils.datasets.distdataset import DistDataset
+from hydragnn.utils.datasets.pickledataset import (
+    SimplePickleWriter,
+    SimplePickleDataset,
+)
+from hydragnn.preprocess.graph_samples_checks_and_updates import gather_deg
 
 try:
-    from hydragnn.utils.adiosdataset import AdiosWriter, AdiosDataset
+    from hydragnn.utils.datasets.adiosdataset import AdiosWriter, AdiosDataset
 except ImportError:
     pass
 
@@ -117,7 +120,7 @@
 
     ##################################################################################################################
     # Always initialize for multi-rank training.
-    comm_size, rank = hydragnn.utils.setup_ddp()
+    comm_size, rank = hydragnn.utils.distributed.setup_ddp()
     ##################################################################################################################
 
     comm = MPI.COMM_WORLD
@@ -130,8 +133,8 @@
     )
 
     log_name = "LJ" if args.log is None else args.log
-    hydragnn.utils.setup_log(log_name)
-    writer = hydragnn.utils.get_summary_writer(log_name)
+    hydragnn.utils.print.setup_log(log_name)
+    writer = hydragnn.utils.model.get_summary_writer(log_name)
 
     log("Command: {0}\n".format(" ".join([x for x in sys.argv])), rank=0)
 
@@ -266,11 +269,13 @@
         trainset, valset, testset, config["NeuralNetwork"]["Training"]["batch_size"]
     )
 
-    config = hydragnn.utils.update_config(config, train_loader, val_loader, test_loader)
+    config = hydragnn.utils.input_config_parsing.update_config(
+        config, train_loader, val_loader, test_loader
+    )
     ## Good to sync with everyone right after DDStore setup
     comm.Barrier()
 
-    hydragnn.utils.save_config(config, log_name)
+    hydragnn.utils.input_config_parsing.save_config(config, log_name)
 
     timer.stop()
 
@@ -278,7 +283,7 @@
         config=config["NeuralNetwork"],
         verbosity=verbosity,
     )
-    model = hydragnn.utils.get_distributed_model(model, verbosity)
+    model = hydragnn.utils.distributed.get_distributed_model(model, verbosity)
 
     learning_rate = config["NeuralNetwork"]["Training"]["Optimizer"]["learning_rate"]
     optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
@@ -286,7 +291,7 @@
         optimizer, mode="min", factor=0.5, patience=5, min_lr=0.00001
     )
 
-    hydragnn.utils.load_existing_model_config(
+    hydragnn.utils.model.load_existing_model_config(
         model, config["NeuralNetwork"]["Training"], optimizer=optimizer
     )
 
@@ -307,8 +312,8 @@
         compute_grad_energy=True,
     )
 
-    hydragnn.utils.save_model(model, optimizer, log_name)
-    hydragnn.utils.print_timers(verbosity)
+    hydragnn.utils.model.save_model(model, optimizer, log_name)
+    hydragnn.utils.profiling_and_tracing.print_timers(verbosity)
 
     if tr.has("GPTLTracer"):
         import gptl4py as gp
diff --git a/examples/alexandria/train.py b/examples/alexandria/train.py
index 15624c03d..7382d2bb8 100644
--- a/examples/alexandria/train.py
+++ b/examples/alexandria/train.py
@@ -15,17 +15,23 @@
 from torch_geometric.transforms import Distance, Spherical, LocalCartesian
 
 import hydragnn
-from hydragnn.utils.time_utils import Timer
+from hydragnn.utils.profiling_and_tracing.time_utils import Timer
 from hydragnn.utils.model import print_model
-from hydragnn.utils.abstractbasedataset import AbstractBaseDataset
-from hydragnn.utils.distdataset import DistDataset
-from hydragnn.utils.pickledataset import SimplePickleWriter, SimplePickleDataset
-from hydragnn.preprocess.utils import gather_deg
-from hydragnn.preprocess.utils import RadiusGraph, RadiusGraphPBC
+from hydragnn.utils.datasets.abstractbasedataset import AbstractBaseDataset
+from hydragnn.utils.datasets.distdataset import DistDataset
+from hydragnn.utils.datasets.pickledataset import (
+    SimplePickleWriter,
+    SimplePickleDataset,
+)
+from hydragnn.preprocess.graph_samples_checks_and_updates import gather_deg
+from hydragnn.preprocess.graph_samples_checks_and_updates import (
+    RadiusGraph,
+    RadiusGraphPBC,
+)
 from hydragnn.preprocess.load_data import split_dataset
 
-import hydragnn.utils.tracer as tr
-from hydragnn.utils.print_utils import iterate_tqdm, log
+import hydragnn.utils.profiling_and_tracing.tracer as tr
+from hydragnn.utils.print.print_utils import iterate_tqdm, log
 
 from generate_dictionaries_pure_elements import (
     generate_dictionary_bulk_energies,
@@ -38,7 +44,7 @@
     pass
 
 import subprocess
-from hydragnn.utils import nsplit
+from hydragnn.utils.distributed import nsplit
 
 
 def info(*args, logtype="info", sep=" "):
@@ -360,7 +366,7 @@ def get(self, idx):
 
     ##################################################################################################################
     # Always initialize for multi-rank training.
-    comm_size, rank = hydragnn.utils.setup_ddp()
+    comm_size, rank = hydragnn.utils.distributed.setup_ddp()
     ##################################################################################################################
 
     comm = MPI.COMM_WORLD
@@ -373,8 +379,8 @@ def get(self, idx):
     )
 
     log_name = "Alexandria" if args.log is None else args.log
-    hydragnn.utils.setup_log(log_name)
-    writer = hydragnn.utils.get_summary_writer(log_name)
+    hydragnn.utils.print.setup_log(log_name)
+    writer = hydragnn.utils.model.get_summary_writer(log_name)
 
     log("Command: {0}\n".format(" ".join([x for x in sys.argv])), rank=0)
 
@@ -507,11 +513,13 @@ def get(self, idx):
         trainset, valset, testset, config["NeuralNetwork"]["Training"]["batch_size"]
     )
 
-    config = hydragnn.utils.update_config(config, train_loader, val_loader, test_loader)
+    config = hydragnn.utils.input_config_parsing.update_config(
+        config, train_loader, val_loader, test_loader
+    )
     ## Good to sync with everyone right after DDStore setup
     comm.Barrier()
 
-    hydragnn.utils.save_config(config, log_name)
+    hydragnn.utils.input_config_parsing.save_config(config, log_name)
 
     timer.stop()
 
@@ -519,7 +527,7 @@ def get(self, idx):
         config=config["NeuralNetwork"],
         verbosity=verbosity,
     )
-    model = hydragnn.utils.get_distributed_model(model, verbosity)
+    model = hydragnn.utils.distributed.get_distributed_model(model, verbosity)
 
     # Print details of neural network architecture
     print_model(model)
@@ -530,7 +538,7 @@ def get(self, idx):
         optimizer, mode="min", factor=0.5, patience=5, min_lr=0.00001
     )
 
-    hydragnn.utils.load_existing_model_config(
+    hydragnn.utils.model.load_existing_model_config(
         model, config["NeuralNetwork"]["Training"], optimizer=optimizer
     )
 
@@ -550,8 +558,8 @@ def get(self, idx):
         create_plots=False,
     )
 
-    hydragnn.utils.save_model(model, optimizer, log_name)
-    hydragnn.utils.print_timers(verbosity)
+    hydragnn.utils.model.save_model(model, optimizer, log_name)
+    hydragnn.utils.profiling_and_tracing.print_timers(verbosity)
 
     if tr.has("GPTLTracer"):
         import gptl4py as gp
diff --git a/examples/ani1_x/train.py b/examples/ani1_x/train.py
index 89831e3b5..100ad56d2 100644
--- a/examples/ani1_x/train.py
+++ b/examples/ani1_x/train.py
@@ -1,4 +1,4 @@
-import os, re, json
+import os, json
 import logging
 import sys
 from mpi4py import MPI
@@ -7,7 +7,6 @@
 import numpy as np
 
 import random
-
 import torch
 
 # FIX random seed
@@ -18,26 +17,31 @@
 from torch_geometric.transforms import Distance, Spherical, LocalCartesian
 
 import hydragnn
-from hydragnn.utils.time_utils import Timer
+from hydragnn.utils.profiling_and_tracing.time_utils import Timer
 from hydragnn.utils.model import print_model
-from hydragnn.utils.abstractbasedataset import AbstractBaseDataset
-from hydragnn.utils.distdataset import DistDataset
-from hydragnn.utils.pickledataset import SimplePickleWriter, SimplePickleDataset
-from hydragnn.preprocess.utils import gather_deg
-from hydragnn.preprocess.utils import RadiusGraph, RadiusGraphPBC
+from hydragnn.utils.datasets.abstractbasedataset import AbstractBaseDataset
+from hydragnn.utils.datasets.distdataset import DistDataset
+from hydragnn.utils.datasets.pickledataset import (
+    SimplePickleWriter,
+    SimplePickleDataset,
+)
+from hydragnn.preprocess.graph_samples_checks_and_updates import gather_deg
+from hydragnn.preprocess.graph_samples_checks_and_updates import (
+    RadiusGraph,
+    RadiusGraphPBC,
+)
 from hydragnn.preprocess.load_data import split_dataset
 
-import hydragnn.utils.tracer as tr
+import hydragnn.utils.profiling_and_tracing.tracer as tr
 
-from hydragnn.utils.print_utils import iterate_tqdm, log
+from hydragnn.utils.print.print_utils import log
 
 try:
     from hydragnn.utils.adiosdataset import AdiosWriter, AdiosDataset
 except ImportError:
     pass
 
-import subprocess
-from hydragnn.utils import nsplit
+from hydragnn.utils.distributed import nsplit
 
 import h5py
 
@@ -242,7 +246,7 @@ def get(self, idx):
 
     ##################################################################################################################
     # Always initialize for multi-rank training.
-    comm_size, rank = hydragnn.utils.setup_ddp()
+    comm_size, rank = hydragnn.utils.distributed.setup_ddp()
     ##################################################################################################################
 
     comm = MPI.COMM_WORLD
@@ -255,8 +259,8 @@ def get(self, idx):
     )
 
     log_name = "ANI1x" if args.log is None else args.log
-    hydragnn.utils.setup_log(log_name)
-    writer = hydragnn.utils.get_summary_writer(log_name)
+    hydragnn.utils.print.print_utils.setup_log(log_name)
+    writer = hydragnn.utils.model.get_summary_writer(log_name)
 
     log("Command: {0}\n".format(" ".join([x for x in sys.argv])), rank=0)
 
@@ -389,11 +393,13 @@ def get(self, idx):
         trainset, valset, testset, config["NeuralNetwork"]["Training"]["batch_size"]
     )
 
-    config = hydragnn.utils.update_config(config, train_loader, val_loader, test_loader)
+    config = hydragnn.utils.input_config_parsing.update_config(
+        config, train_loader, val_loader, test_loader
+    )
     ## Good to sync with everyone right after DDStore setup
     comm.Barrier()
 
-    hydragnn.utils.save_config(config, log_name)
+    hydragnn.utils.input_config_parsing.save_config(config, log_name)
 
     timer.stop()
 
@@ -401,7 +407,7 @@ def get(self, idx):
         config=config["NeuralNetwork"],
         verbosity=verbosity,
     )
-    model = hydragnn.utils.get_distributed_model(model, verbosity)
+    model = hydragnn.utils.distributed.get_distributed_model(model, verbosity)
 
     # Print details of neural network architecture
     print_model(model)
@@ -412,7 +418,7 @@ def get(self, idx):
         optimizer, mode="min", factor=0.5, patience=5, min_lr=0.00001
     )
 
-    hydragnn.utils.load_existing_model_config(
+    hydragnn.utils.model.load_existing_model_config(
         model, config["NeuralNetwork"]["Training"], optimizer=optimizer
     )
 
@@ -432,8 +438,8 @@ def get(self, idx):
         create_plots=False,
     )
 
-    hydragnn.utils.save_model(model, optimizer, log_name)
-    hydragnn.utils.print_timers(verbosity)
+    hydragnn.utils.model.save_model(model, optimizer, log_name)
+    hydragnn.utils.profiling_and_tracing.print_timers(verbosity)
 
     if tr.has("GPTLTracer"):
         import gptl4py as gp
diff --git a/examples/csce/train_gap.py b/examples/csce/train_gap.py
index d802ad50d..f4d2cacba 100644
--- a/examples/csce/train_gap.py
+++ b/examples/csce/train_gap.py
@@ -9,21 +9,26 @@
 
 import logging
 import sys
+
 from mpi4py import MPI
+
 import argparse
 
 import hydragnn
-from hydragnn.utils.print_utils import print_distributed, iterate_tqdm, log
-from hydragnn.utils.time_utils import Timer
-from hydragnn.utils.distdataset import DistDataset
-from hydragnn.utils.pickledataset import SimplePickleWriter, SimplePickleDataset
-from hydragnn.utils.smiles_utils import (
+from hydragnn.utils.print.print_utils import print_distributed, iterate_tqdm, log
+from hydragnn.utils.profiling_and_tracing.time_utils import Timer
+from hydragnn.utils.datasets.distdataset import DistDataset
+from hydragnn.utils.datasets.pickledataset import (
+    SimplePickleWriter,
+    SimplePickleDataset,
+)
+from hydragnn.utils.descriptors_and_embeddings.smiles_utils import (
     get_node_attribute_name,
     generate_graphdata_from_smilestr,
 )
-from hydragnn.preprocess.utils import gather_deg
-from hydragnn.utils import nsplit
-import hydragnn.utils.tracer as tr
+from hydragnn.preprocess.graph_samples_checks_and_updates import gather_deg
+from hydragnn.utils.distributed import nsplit
+import hydragnn.utils.profiling_and_tracing.tracer as tr
 
 import numpy as np
 
@@ -228,7 +233,7 @@ def __getitem__(self, idx):
     ) = get_node_attribute_name(csce_node_types)
     ##################################################################################################################
     # Always initialize for multi-rank training.
-    comm_size, rank = hydragnn.utils.setup_ddp()
+    comm_size, rank = hydragnn.utils.distributed.setup_ddp()
     ##################################################################################################################
 
     comm = MPI.COMM_WORLD
@@ -243,8 +248,8 @@ def __getitem__(self, idx):
     log_name = "csce_" + inputfilesubstr + "_eV_fullx"
     if args.log is not None:
         log_name = args.log
-    hydragnn.utils.setup_log(log_name)
-    writer = hydragnn.utils.get_summary_writer(log_name)
+    hydragnn.utils.print.print_utils.setup_log(log_name)
+    writer = hydragnn.utils.model.get_summary_writer(log_name)
 
     log("Command: {0}\n".format(" ".join([x for x in sys.argv])), rank=0)
 
@@ -343,9 +348,7 @@ def __getitem__(self, idx):
             os.environ["HYDRAGNN_USE_ddstore"] = "1"
 
         opt = {"preload": False, "shmem": shmem, "ddstore": ddstore}
-        fname = fname = os.path.join(
-            os.path.dirname(__file__), "dataset", "csce_gap.bp"
-        )
+        fname = os.path.join(os.path.dirname(__file__), "dataset", "csce_gap.bp")
         trainset = AdiosDataset(fname, "trainset", comm, **opt)
         valset = AdiosDataset(fname, "valset", comm)
         testset = AdiosDataset(fname, "testset", comm)
@@ -385,10 +388,12 @@ def __getitem__(self, idx):
     )
     comm.Barrier()
 
-    config = hydragnn.utils.update_config(config, train_loader, val_loader, test_loader)
+    config = hydragnn.utils.input_config_parsing.update_config(
+        config, train_loader, val_loader, test_loader
+    )
     comm.Barrier()
 
-    hydragnn.utils.save_config(config, log_name)
+    hydragnn.utils.input_config_parsing.save_config(config, log_name)
     comm.Barrier()
 
     timer.stop()
@@ -397,7 +402,7 @@ def __getitem__(self, idx):
         config=config["NeuralNetwork"],
         verbosity=verbosity,
     )
-    model = hydragnn.utils.get_distributed_model(model, verbosity)
+    model = hydragnn.utils.distributed.get_distributed_model(model, verbosity)
 
     learning_rate = config["NeuralNetwork"]["Training"]["Optimizer"]["learning_rate"]
     optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
@@ -405,7 +410,7 @@ def __getitem__(self, idx):
         optimizer, mode="min", factor=0.5, patience=5, min_lr=0.00001
     )
 
-    hydragnn.utils.load_existing_model_config(
+    hydragnn.utils.model.load_existing_model_config(
         model, config["NeuralNetwork"]["Training"], optimizer=optimizer
     )
 
@@ -425,8 +430,8 @@ def __getitem__(self, idx):
         create_plots=False,
     )
 
-    hydragnn.utils.save_model(model, optimizer, log_name)
-    hydragnn.utils.print_timers(verbosity)
+    hydragnn.utils.model.save_model(model, optimizer, log_name)
+    hydragnn.utils.profiling_and_tracing.time_utils.print_timers(verbosity)
 
     if args.mae:
         import matplotlib.pyplot as plt
diff --git a/examples/dftb_uv_spectrum/train_discrete_uv_spectrum.py b/examples/dftb_uv_spectrum/train_discrete_uv_spectrum.py
index b81fd5b9b..0f6484911 100644
--- a/examples/dftb_uv_spectrum/train_discrete_uv_spectrum.py
+++ b/examples/dftb_uv_spectrum/train_discrete_uv_spectrum.py
@@ -5,31 +5,29 @@
 
 import os, json
 import random
-import pickle
 
 import logging
 import sys
-from tqdm import tqdm
 from mpi4py import MPI
-from itertools import chain
 import argparse
-import time
 
 from rdkit.Chem.rdmolfiles import MolFromPDBFile
 
 import hydragnn
-from hydragnn.utils.print_utils import print_distributed, iterate_tqdm, log
-from hydragnn.utils.time_utils import Timer
-from hydragnn.utils.pickledataset import SimplePickleDataset
-from hydragnn.utils.smiles_utils import (
+from hydragnn.utils.print.print_utils import print_distributed, iterate_tqdm, log
+from hydragnn.utils.profiling_and_tracing.time_utils import Timer
+from hydragnn.utils.descriptors_and_embeddings.smiles_utils import (
     get_node_attribute_name,
     generate_graphdata_from_rdkit_molecule,
 )
 from hydragnn.utils.distributed import get_device
 from hydragnn.preprocess.load_data import split_dataset
-from hydragnn.utils.distdataset import DistDataset
-from hydragnn.utils.pickledataset import SimplePickleWriter, SimplePickleDataset
-from hydragnn.preprocess.utils import gather_deg
+from hydragnn.utils.datasets.distdataset import DistDataset
+from hydragnn.utils.datasets.pickledataset import (
+    SimplePickleWriter,
+    SimplePickleDataset,
+)
+from hydragnn.preprocess.graph_samples_checks_and_updates import gather_deg
 
 import numpy as np
 
@@ -44,9 +42,9 @@
 
 import warnings
 
-from hydragnn.utils import nsplit
-import hydragnn.utils.tracer as tr
-from hydragnn.utils.abstractbasedataset import AbstractBaseDataset
+from hydragnn.utils.distributed import nsplit
+import hydragnn.utils.profiling_and_tracing.tracer as tr
+from hydragnn.utils.datasets.abstractbasedataset import AbstractBaseDataset
 
 # FIXME: this works fine for now because we train on GDB-9 molecules
 # for larger chemical spaces, the following atom representation has to be properly expanded
@@ -58,6 +56,7 @@ def info(*args, logtype="info", sep=" "):
 
 
 def dftb_to_graph(moldir, dftb_node_types, var_config):
+    pdb_filename = os.path.join(moldir, "smiles.pdb")
     pdb_filename = os.path.join(moldir, "smiles.pdb")
     mol = MolFromPDBFile(
         pdb_filename, sanitize=False, proximityBonding=True, removeHs=True
diff --git a/examples/dftb_uv_spectrum/train_smooth_uv_spectrum.py b/examples/dftb_uv_spectrum/train_smooth_uv_spectrum.py
index 5af612c4a..4233d9c05 100644
--- a/examples/dftb_uv_spectrum/train_smooth_uv_spectrum.py
+++ b/examples/dftb_uv_spectrum/train_smooth_uv_spectrum.py
@@ -5,33 +5,31 @@
 
 import os, json
 import random
-import pickle
 
 import logging
 import sys
-from tqdm import tqdm
 from mpi4py import MPI
-from itertools import chain
 import argparse
-import time
 
 from rdkit.Chem.rdmolfiles import MolFromPDBFile
 
 import hydragnn
-from hydragnn.utils.print_utils import print_distributed, iterate_tqdm, log
-from hydragnn.utils.time_utils import Timer
+from hydragnn.utils.print.print_utils import print_distributed, iterate_tqdm, log
+from hydragnn.utils.profiling_and_tracing.time_utils import Timer
 
 # from hydragnn.utils.adiosdataset import AdiosWriter, AdiosDataset
-from hydragnn.utils.pickledataset import SimplePickleDataset
-from hydragnn.utils.smiles_utils import (
+from hydragnn.utils.descriptors_and_embeddings.smiles_utils import (
     get_node_attribute_name,
     generate_graphdata_from_rdkit_molecule,
 )
 from hydragnn.utils.distributed import get_device
 from hydragnn.preprocess.load_data import split_dataset
-from hydragnn.utils.distdataset import DistDataset
-from hydragnn.utils.pickledataset import SimplePickleWriter, SimplePickleDataset
-from hydragnn.preprocess.utils import gather_deg
+from hydragnn.utils.datasets.distdataset import DistDataset
+from hydragnn.utils.datasets.pickledataset import (
+    SimplePickleWriter,
+    SimplePickleDataset,
+)
+from hydragnn.preprocess.graph_samples_checks_and_updates import gather_deg
 
 import numpy as np
 
@@ -40,12 +38,11 @@
 except ImportError:
     pass
 
-import torch_geometric.data
 import torch
 import torch.distributed as dist
 
-from hydragnn.utils import nsplit
-import hydragnn.utils.tracer as tr
+from hydragnn.utils.distributed import nsplit
+import hydragnn.utils.profiling_and_tracing.tracer as tr
 
 # FIXME: this works fine for now because we train on GDB-9 molecules
 # for larger chemical spaces, the following atom representation has to be properly expanded
@@ -56,7 +53,7 @@ def info(*args, logtype="info", sep=" "):
     getattr(logging, logtype)(sep.join(map(str, args)))
 
 
-from hydragnn.utils.abstractbasedataset import AbstractBaseDataset
+from hydragnn.utils.datasets.abstractbasedataset import AbstractBaseDataset
 
 
 def dftb_to_graph(moldir, dftb_node_types, var_config):
diff --git a/examples/eam/eam.py b/examples/eam/eam.py
index 64a8b804e..4de2127b6 100644
--- a/examples/eam/eam.py
+++ b/examples/eam/eam.py
@@ -5,13 +5,15 @@
 import argparse
 
 import hydragnn
-from hydragnn.utils.time_utils import Timer
-from hydragnn.utils.config_utils import get_log_name_config
+from hydragnn.utils.profiling_and_tracing.time_utils import Timer
+from hydragnn.utils.input_config_parsing.config_utils import get_log_name_config
 from hydragnn.utils.model import print_model
-from hydragnn.utils.cfgdataset import CFGDataset
-from hydragnn.utils.serializeddataset import SerializedWriter, SerializedDataset
+from hydragnn.utils.datasets.cfgdataset import CFGDataset
+from hydragnn.utils.datasets.serializeddataset import (
+    SerializedWriter,
+    SerializedDataset,
+)
 from hydragnn.preprocess.load_data import split_dataset
-from hydragnn.utils.print_utils import log
 
 try:
     from hydragnn.utils.adiosdataset import AdiosWriter, AdiosDataset
@@ -64,10 +66,10 @@ def info(*args, logtype="info", sep=" "):
     input_filename = os.path.join(dirpwd, args.inputfile)
     with open(input_filename, "r") as f:
         config = json.load(f)
-    hydragnn.utils.setup_log(get_log_name_config(config))
+    hydragnn.utils.print.setup_log(get_log_name_config(config))
     ##################################################################################################################
     # Always initialize for multi-rank training.
-    comm_size, rank = hydragnn.utils.setup_ddp()
+    comm_size, rank = hydragnn.utils.distributed.setup_ddp()
     ##################################################################################################################
     comm = MPI.COMM_WORLD
     ## Set up logging
@@ -170,7 +172,9 @@ def info(*args, logtype="info", sep=" "):
     )
     timer.stop()
 
-    config = hydragnn.utils.update_config(config, train_loader, val_loader, test_loader)
+    config = hydragnn.utils.input_config_parsing.update_config(
+        config, train_loader, val_loader, test_loader
+    )
     config["NeuralNetwork"]["Variables_of_interest"].pop("minmax_node_feature", None)
     config["NeuralNetwork"]["Variables_of_interest"].pop("minmax_graph_feature", None)
 
@@ -183,7 +187,7 @@ def info(*args, logtype="info", sep=" "):
         print_model(model)
     comm.Barrier()
 
-    model = hydragnn.utils.get_distributed_model(model, verbosity)
+    model = hydragnn.utils.distributed.get_distributed_model(model, verbosity)
 
     learning_rate = config["NeuralNetwork"]["Training"]["Optimizer"]["learning_rate"]
     optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
@@ -192,12 +196,12 @@ def info(*args, logtype="info", sep=" "):
     )
 
     log_name = get_log_name_config(config)
-    writer = hydragnn.utils.get_summary_writer(log_name)
+    writer = hydragnn.utils.model.get_summary_writer(log_name)
 
     if dist.is_initialized():
         dist.barrier()
 
-    hydragnn.utils.save_config(config, log_name)
+    hydragnn.utils.input_config_parsing.save_config(config, log_name)
 
     hydragnn.train.train_validate_test(
         model,
@@ -213,7 +217,7 @@ def info(*args, logtype="info", sep=" "):
         create_plots=True,
     )
 
-    hydragnn.utils.save_model(model, optimizer, log_name)
-    hydragnn.utils.print_timers(verbosity)
+    hydragnn.utils.model.save_model(model, optimizer, log_name)
+    hydragnn.utils.profiling_and_tracing.time_utils.print_timers(verbosity)
 
     sys.exit(0)
diff --git a/examples/ising_model/train_ising.py b/examples/ising_model/train_ising.py
index 892b08824..3b8c859bd 100644
--- a/examples/ising_model/train_ising.py
+++ b/examples/ising_model/train_ising.py
@@ -11,20 +11,23 @@
 import argparse
 
 import hydragnn
-from hydragnn.utils.print_utils import print_distributed, iterate_tqdm, log
-from hydragnn.utils.time_utils import Timer
-from hydragnn.utils.config_utils import get_log_name_config
+from hydragnn.utils.print.print_utils import print_distributed, iterate_tqdm, log
+from hydragnn.utils.profiling_and_tracing.time_utils import Timer
+from hydragnn.utils.input_config_parsing.config_utils import get_log_name_config
 from hydragnn.preprocess.load_data import split_dataset
 from hydragnn.utils.model import print_model
-from hydragnn.utils.lsmsdataset import LSMSDataset
-from hydragnn.utils.distdataset import DistDataset
-from hydragnn.utils.pickledataset import SimplePickleWriter, SimplePickleDataset
-from hydragnn.preprocess.utils import gather_deg
+from hydragnn.utils.datasets.lsmsdataset import LSMSDataset
+from hydragnn.utils.datasets.distdataset import DistDataset
+from hydragnn.utils.datasets.pickledataset import (
+    SimplePickleWriter,
+    SimplePickleDataset,
+)
+from hydragnn.preprocess.graph_samples_checks_and_updates import gather_deg
 
 import numpy as np
 
 try:
-    from hydragnn.utils.adiosdataset import AdiosWriter, AdiosDataset
+    from hydragnn.utils.datasets.adiosdataset import AdiosWriter, AdiosDataset
 except ImportError:
     pass
 
@@ -45,8 +48,8 @@
 
 from create_configurations import E_dimensionless
 
-import hydragnn.utils.tracer as tr
-from hydragnn.utils import nsplit
+import hydragnn.utils.profiling_and_tracing.tracer as tr
+from hydragnn.utils.distributed import nsplit
 
 
 def write_to_file(total_energy, atomic_features, count_config, dir, prefix):
@@ -189,10 +192,10 @@ def info(*args, logtype="info", sep=" "):
     log_name = get_log_name_config(config)
     if args.log is not None:
         log_name = args.log
-    hydragnn.utils.setup_log(log_name)
+    hydragnn.utils.print.print_utils.setup_log(log_name)
     ##################################################################################################################
     # Always initialize for multi-rank training.
-    comm_size, rank = hydragnn.utils.setup_ddp()
+    comm_size, rank = hydragnn.utils.distributed.setup_ddp()
     ##################################################################################################################
 
     comm = MPI.COMM_WORLD
@@ -364,7 +367,9 @@ def info(*args, logtype="info", sep=" "):
     config["NeuralNetwork"]["Variables_of_interest"][
         "minmax_graph_feature"
     ] = trainset.minmax_graph_feature
-    config = hydragnn.utils.update_config(config, train_loader, val_loader, test_loader)
+    config = hydragnn.utils.input_config_parsing.update_config(
+        config, train_loader, val_loader, test_loader
+    )
     del config["NeuralNetwork"]["Variables_of_interest"]["minmax_node_feature"]
     del config["NeuralNetwork"]["Variables_of_interest"]["minmax_graph_feature"]
     ## Good to sync with everyone right after DDStore setup
@@ -379,7 +384,7 @@ def info(*args, logtype="info", sep=" "):
         print_model(model)
     comm.Barrier()
 
-    model = hydragnn.utils.get_distributed_model(model, verbosity)
+    model = hydragnn.utils.distributed.get_distributed_model(model, verbosity)
 
     learning_rate = config["NeuralNetwork"]["Training"]["Optimizer"]["learning_rate"]
     optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
@@ -387,7 +392,7 @@ def info(*args, logtype="info", sep=" "):
         optimizer, mode="min", factor=0.5, patience=5, min_lr=0.00001
     )
 
-    writer = hydragnn.utils.get_summary_writer(log_name)
+    writer = hydragnn.utils.model.get_summary_writer(log_name)
 
     if dist.is_initialized():
         dist.barrier()
@@ -407,8 +412,8 @@ def info(*args, logtype="info", sep=" "):
         verbosity,
     )
 
-    hydragnn.utils.save_model(model, optimizer, log_name)
-    hydragnn.utils.print_timers(verbosity)
+    hydragnn.utils.model.save_model(model, optimizer, log_name)
+    hydragnn.utils.profiling_and_tracing.print_timers(verbosity)
 
     if tr.has("GPTLTracer"):
         import gptl4py as gp
diff --git a/examples/lsms/lsms.py b/examples/lsms/lsms.py
index 1c97ba33f..8a64e353c 100644
--- a/examples/lsms/lsms.py
+++ b/examples/lsms/lsms.py
@@ -5,11 +5,14 @@
 import argparse
 
 import hydragnn
-from hydragnn.utils.time_utils import Timer
-from hydragnn.utils.config_utils import get_log_name_config
+from hydragnn.utils.profiling_and_tracing.time_utils import Timer
+from hydragnn.utils.input_config_parsing.config_utils import get_log_name_config
 from hydragnn.utils.model import print_model
-from hydragnn.utils.lsmsdataset import LSMSDataset
-from hydragnn.utils.serializeddataset import SerializedWriter, SerializedDataset
+from hydragnn.utils.datasets.lsmsdataset import LSMSDataset
+from hydragnn.utils.datasets.serializeddataset import (
+    SerializedWriter,
+    SerializedDataset,
+)
 from hydragnn.preprocess.load_data import split_dataset
 
 try:
@@ -65,10 +68,10 @@ def info(*args, logtype="info", sep=" "):
     input_filename = os.path.join(dirpwd, args.inputfile)
     with open(input_filename, "r") as f:
         config = json.load(f)
-    hydragnn.utils.setup_log(get_log_name_config(config))
+    hydragnn.utils.print.setup_log(get_log_name_config(config))
     ##################################################################################################################
     # Always initialize for multi-rank training.
-    comm_size, rank = hydragnn.utils.setup_ddp()
+    comm_size, rank = hydragnn.utils.distributed.setup_ddp()
     ##################################################################################################################
     comm = MPI.COMM_WORLD
     ## Set up logging
@@ -172,7 +175,9 @@ def info(*args, logtype="info", sep=" "):
     )
     timer.stop()
 
-    config = hydragnn.utils.update_config(config, train_loader, val_loader, test_loader)
+    config = hydragnn.utils.input_config_parsing.update_config(
+        config, train_loader, val_loader, test_loader
+    )
     config["NeuralNetwork"]["Variables_of_interest"].pop("minmax_node_feature", None)
     config["NeuralNetwork"]["Variables_of_interest"].pop("minmax_graph_feature", None)
 
@@ -185,7 +190,7 @@ def info(*args, logtype="info", sep=" "):
         print_model(model)
     comm.Barrier()
 
-    model = hydragnn.utils.get_distributed_model(model, verbosity)
+    model = hydragnn.utils.distributed.get_distributed_model(model, verbosity)
 
     learning_rate = config["NeuralNetwork"]["Training"]["Optimizer"]["learning_rate"]
     optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
@@ -194,12 +199,12 @@ def info(*args, logtype="info", sep=" "):
     )
 
     log_name = get_log_name_config(config)
-    writer = hydragnn.utils.get_summary_writer(log_name)
+    writer = hydragnn.utils.model.get_summary_writer(log_name)
 
     if dist.is_initialized():
         dist.barrier()
 
-    hydragnn.utils.save_config(config, log_name)
+    hydragnn.utils.input_config_parsing.save_config(config, log_name)
 
     hydragnn.train.train_validate_test(
         model,
@@ -215,7 +220,7 @@ def info(*args, logtype="info", sep=" "):
         create_plots=True,
     )
 
-    hydragnn.utils.save_model(model, optimizer, log_name)
-    hydragnn.utils.print_timers(verbosity)
+    hydragnn.utils.model.save_model(model, optimizer, log_name)
+    hydragnn.utils.profiling_and_tracing.print_timers(verbosity)
 
     sys.exit(0)
diff --git a/examples/md17/md17.py b/examples/md17/md17.py
index 7e71b872e..83695df9c 100644
--- a/examples/md17/md17.py
+++ b/examples/md17/md17.py
@@ -52,15 +52,15 @@ def md17_pre_filter(data):
 
 log_name = "md17_test"
 # Enable print to log file.
-hydragnn.utils.setup_log(log_name)
+hydragnn.utils.print.print_utils.setup_log(log_name)
 
-# Use built-in torch_geometric dataset.
+# Use built-in torch_geometric datasets.
 # Filter function above used to run quick example.
 # NOTE: data is moved to the device in the pre-transform.
 # NOTE: transforms/filters will NOT be re-run unless the qm9/processed/ directory is removed.
 compute_edges = hydragnn.preprocess.get_radius_graph_config(arch_config)
 
-# Fix for MD17 dataset
+# Fix for MD17 datasets
 torch_geometric.datasets.MD17.file_names["uracil"] = "md17_uracil.npz"
 
 dataset = torch_geometric.datasets.MD17(
@@ -76,13 +76,15 @@ def md17_pre_filter(data):
     train, val, test, config["NeuralNetwork"]["Training"]["batch_size"]
 )
 
-config = hydragnn.utils.update_config(config, train_loader, val_loader, test_loader)
+config = hydragnn.utils.input_config_parsing.update_config(
+    config, train_loader, val_loader, test_loader
+)
 
 model = hydragnn.models.create_model_config(
     config=config["NeuralNetwork"],
     verbosity=verbosity,
 )
-model = hydragnn.utils.get_distributed_model(model, verbosity)
+model = hydragnn.utils.distributed.get_distributed_model(model, verbosity)
 
 learning_rate = config["NeuralNetwork"]["Training"]["Optimizer"]["learning_rate"]
 optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
@@ -90,9 +92,9 @@ def md17_pre_filter(data):
     optimizer, mode="min", factor=0.5, patience=5, min_lr=0.00001
 )
 
-# Run training with the given model and qm9 dataset.
-writer = hydragnn.utils.get_summary_writer(log_name)
-hydragnn.utils.save_config(config, log_name)
+# Run training with the given model and md17 dataset.
+writer = hydragnn.utils.model.model.get_summary_writer(log_name)
+hydragnn.utils.input_config_parsing.save_config(config, log_name)
 
 hydragnn.train.train_validate_test(
     model,
diff --git a/examples/mptrj/train.py b/examples/mptrj/train.py
index 36f9ba821..8dee83db6 100644
--- a/examples/mptrj/train.py
+++ b/examples/mptrj/train.py
@@ -1,4 +1,4 @@
-import os, re, json
+import os, json
 import logging
 import sys
 from mpi4py import MPI
@@ -17,18 +17,24 @@
 from torch_geometric.transforms import Distance, Spherical, LocalCartesian
 
 import hydragnn
-from hydragnn.utils.time_utils import Timer
+from hydragnn.utils.profiling_and_tracing.time_utils import Timer
 from hydragnn.utils.model import print_model
-from hydragnn.utils.abstractbasedataset import AbstractBaseDataset
-from hydragnn.utils.distdataset import DistDataset
-from hydragnn.utils.pickledataset import SimplePickleWriter, SimplePickleDataset
-from hydragnn.preprocess.utils import gather_deg
-from hydragnn.preprocess.utils import RadiusGraph, RadiusGraphPBC
+from hydragnn.utils.datasets.abstractbasedataset import AbstractBaseDataset
+from hydragnn.utils.datasets.distdataset import DistDataset
+from hydragnn.utils.datasets.pickledataset import (
+    SimplePickleWriter,
+    SimplePickleDataset,
+)
+from hydragnn.preprocess.graph_samples_checks_and_updates import gather_deg
+from hydragnn.preprocess.graph_samples_checks_and_updates import (
+    RadiusGraph,
+    RadiusGraphPBC,
+)
 from hydragnn.preprocess.load_data import split_dataset
 
-import hydragnn.utils.tracer as tr
+import hydragnn.utils.profiling_and_tracing.tracer as tr
 
-from hydragnn.utils.print_utils import iterate_tqdm, log
+from hydragnn.utils.print.print_utils import iterate_tqdm, log
 
 from jarvis.db.jsonutils import loadjson, dumpjson
 from pymatgen.core.structure import Structure
@@ -42,8 +48,7 @@
 except ImportError:
     pass
 
-import subprocess
-from hydragnn.utils import nsplit
+from hydragnn.utils.distributed import nsplit
 
 
 def info(*args, logtype="info", sep=" "):
@@ -256,7 +261,7 @@ def get(self, idx):
 
     ##################################################################################################################
     # Always initialize for multi-rank training.
-    comm_size, rank = hydragnn.utils.setup_ddp()
+    comm_size, rank = hydragnn.utils.distributed.setup_ddp()
     ##################################################################################################################
 
     comm = MPI.COMM_WORLD
@@ -269,8 +274,8 @@ def get(self, idx):
     )
 
     log_name = "MPTrj" if args.log is None else args.log
-    hydragnn.utils.setup_log(log_name)
-    writer = hydragnn.utils.get_summary_writer(log_name)
+    hydragnn.utils.print.print_utils.setup_log(log_name)
+    writer = hydragnn.utils.model.get_summary_writer(log_name)
 
     log("Command: {0}\n".format(" ".join([x for x in sys.argv])), rank=0)
 
@@ -404,11 +409,13 @@ def get(self, idx):
         trainset, valset, testset, config["NeuralNetwork"]["Training"]["batch_size"]
     )
 
-    config = hydragnn.utils.update_config(config, train_loader, val_loader, test_loader)
+    config = hydragnn.utils.input_config_parsing.update_config(
+        config, train_loader, val_loader, test_loader
+    )
     ## Good to sync with everyone right after DDStore setup
     comm.Barrier()
 
-    hydragnn.utils.save_config(config, log_name)
+    hydragnn.utils.input_config_parsing.save_config(config, log_name)
 
     timer.stop()
 
@@ -416,7 +423,7 @@ def get(self, idx):
         config=config["NeuralNetwork"],
         verbosity=verbosity,
     )
-    model = hydragnn.utils.get_distributed_model(model, verbosity)
+    model = hydragnn.utils.distributed.get_distributed_model(model, verbosity)
 
     # Print details of neural network architecture
     print_model(model)
@@ -427,7 +434,7 @@ def get(self, idx):
         optimizer, mode="min", factor=0.5, patience=5, min_lr=0.00001
     )
 
-    hydragnn.utils.load_existing_model_config(
+    hydragnn.utils.model.load_existing_model_config(
         model, config["NeuralNetwork"]["Training"], optimizer=optimizer
     )
 
@@ -447,8 +454,8 @@ def get(self, idx):
         create_plots=False,
     )
 
-    hydragnn.utils.save_model(model, optimizer, log_name)
-    hydragnn.utils.print_timers(verbosity)
+    hydragnn.utils.model.save_model(model, optimizer, log_name)
+    hydragnn.utils.profiling_and_tracing.print_timers(verbosity)
 
     if tr.has("GPTLTracer"):
         import gptl4py as gp
diff --git a/examples/multidataset/energy_linear_regression.py b/examples/multidataset/energy_linear_regression.py
index e65b5e993..0b81c578b 100644
--- a/examples/multidataset/energy_linear_regression.py
+++ b/examples/multidataset/energy_linear_regression.py
@@ -7,20 +7,25 @@
 import numpy as np
 
 import hydragnn
-from hydragnn.utils import nsplit
-from hydragnn.utils.adiosdataset import AdiosWriter, AdiosDataset
+from hydragnn.utils.distributed import nsplit
+from hydragnn.utils.datasets.adiosdataset import AdiosWriter, AdiosDataset
 from tqdm import tqdm
-from mpi_list import Context, DFM
+
+# This import requires having installed the package mpi_list
+try:
+    from mpi_list import Context, DFM
+except ImportError:
+    print("mpi_list requires having installed: https://github.com/frobnitzem/mpi_list")
 
 
 def subset(i):
-    # sz = len(dataset)
+    # sz = len(datasets)
     # chunk = sz // C.procs
     # left  = sz % C.procs
     # a = i*chunk     + min(i, left)
     # b = (i+1)*chunk + min(i+1, left)
     # print(f"Rank {i}/{C.procs} converting subset [{a},{b})")
-    # return np.array([np.array(x) for x in dataset[a:b]["image"]])
+    # return np.array([np.array(x) for x in datasets[a:b]["image"]])
     return np.random.random((100, 4))
 
 
@@ -101,7 +106,7 @@ def solve_least_squares_svd(A, b):
     comm_rank = comm.Get_rank()
     comm_size = comm.Get_size()
 
-    fname = os.path.join(os.path.dirname(__file__), "./dataset/%s.bp" % args.modelname)
+    fname = os.path.join(os.path.dirname(__file__), "./datasets/%s.bp" % args.modelname)
     print("fname:", fname)
     trainset = AdiosDataset(
         fname,
@@ -123,7 +128,7 @@ def solve_least_squares_svd(A, b):
     )
     pna_deg = trainset.pna_deg
 
-    ## Iterate over local dataset
+    ## Iterate over local datasets
     energy_list = list()
     feature_list = list()
     for dataset in [trainset, valset, testset]:
@@ -205,7 +210,7 @@ def solve_least_squares_svd(A, b):
 
     ## Writing
     fname = os.path.join(
-        os.path.dirname(__file__), "./dataset/%s-v2.bp" % args.modelname
+        os.path.dirname(__file__), "./datasets/%s-v2.bp" % args.modelname
     )
     if comm_rank == 0:
         print("Saving:", fname)
diff --git a/examples/multidataset/train.py b/examples/multidataset/train.py
index 370a58ddf..adc008f01 100644
--- a/examples/multidataset/train.py
+++ b/examples/multidataset/train.py
@@ -13,15 +13,15 @@
 import numpy as np
 
 import hydragnn
-from hydragnn.utils.time_utils import Timer
+from hydragnn.utils.profiling_and_tracing.time_utils import Timer
 from hydragnn.utils.model import print_model
-from hydragnn.utils.distdataset import DistDataset
-from hydragnn.utils.pickledataset import SimplePickleDataset
+from hydragnn.utils.datasets.distdataset import DistDataset
+from hydragnn.utils.datasets.pickledataset import SimplePickleDataset
 
-import hydragnn.utils.tracer as tr
+import hydragnn.utils.profiling_and_tracing.tracer as tr
 
-from hydragnn.utils.print_utils import log, log0
-from hydragnn.utils import nsplit
+from hydragnn.utils.print.print_utils import log, log0
+from hydragnn.utils.distributed import nsplit
 
 try:
     from hydragnn.utils.adiosdataset import AdiosDataset
@@ -122,7 +122,7 @@ def info(*args, logtype="info", sep=" "):
 
     ##################################################################################################################
     # Always initialize for multi-rank training.
-    comm_size, rank = hydragnn.utils.setup_ddp()
+    comm_size, rank = hydragnn.utils.distributed.setup_ddp()
     ##################################################################################################################
 
     comm = MPI.COMM_WORLD
@@ -135,8 +135,8 @@ def info(*args, logtype="info", sep=" "):
     )
 
     log_name = "GFM" if args.log is None else args.log
-    hydragnn.utils.setup_log(log_name)
-    writer = hydragnn.utils.get_summary_writer(log_name)
+    hydragnn.utils.print.setup_log(log_name)
+    writer = hydragnn.utils.model.get_summary_writer(log_name)
 
     log("Command: {0}\n".format(" ".join([x for x in sys.argv])), rank=0)
 
@@ -187,7 +187,7 @@ def info(*args, logtype="info", sep=" "):
             trainset.pna_deg = pna_deg
     elif args.format == "multi":
         ## Reading multiple datasets, which requires the following arguments:
-        ## --multi_model_list: the list datasets/model names
+        ## --multi_model_list: the list dataset/model names
         modellist = args.multi_model_list.split(",")
         if rank == 0:
             ndata_list = list()
@@ -346,11 +346,13 @@ def info(*args, logtype="info", sep=" "):
         test_sampler_shuffle=False,
     )
 
-    config = hydragnn.utils.update_config(config, train_loader, val_loader, test_loader)
+    config = hydragnn.utils.input_config_parsing.update_config(
+        config, train_loader, val_loader, test_loader
+    )
     ## Good to sync with everyone right after DDStore setup
     comm.Barrier()
 
-    hydragnn.utils.save_config(config, log_name)
+    hydragnn.utils.input_config_parsing.save_config(config, log_name)
 
     timer.stop()
 
@@ -358,7 +360,7 @@ def info(*args, logtype="info", sep=" "):
         config=config["NeuralNetwork"],
         verbosity=verbosity,
     )
-    model = hydragnn.utils.get_distributed_model(model, verbosity)
+    model = hydragnn.utils.distributed.get_distributed_model(model, verbosity)
 
     # Print details of neural network architecture
     print_model(model)
@@ -369,7 +371,7 @@ def info(*args, logtype="info", sep=" "):
         optimizer, mode="min", factor=0.5, patience=5, min_lr=0.00001
     )
 
-    hydragnn.utils.load_existing_model_config(
+    hydragnn.utils.model.load_existing_model_config(
         model, config["NeuralNetwork"]["Training"], optimizer=optimizer
     )
 
@@ -389,8 +391,8 @@ def info(*args, logtype="info", sep=" "):
         create_plots=False,
     )
 
-    hydragnn.utils.save_model(model, optimizer, log_name)
-    hydragnn.utils.print_timers(verbosity)
+    hydragnn.utils.model.save_model(model, optimizer, log_name)
+    hydragnn.utils.profiling_and_tracing.print_timers(verbosity)
 
     if tr.has("GPTLTracer"):
         import gptl4py as gp
diff --git a/examples/multidataset_hpo/gfm.py b/examples/multidataset_hpo/gfm.py
index 4228fb39f..daf2bc7f0 100644
--- a/examples/multidataset_hpo/gfm.py
+++ b/examples/multidataset_hpo/gfm.py
@@ -8,15 +8,15 @@
 import numpy as np
 
 import hydragnn
-from hydragnn.utils.time_utils import Timer
+from hydragnn.utils.profiling_and_tracing.time_utils import Timer
 from hydragnn.utils.model import print_model
-from hydragnn.utils.distdataset import DistDataset
-from hydragnn.utils.pickledataset import SimplePickleDataset
+from hydragnn.utils.datasets.distdataset import DistDataset
+from hydragnn.utils.datasets.pickledataset import SimplePickleDataset
 
-import hydragnn.utils.tracer as tr
+import hydragnn.utils.profiling_and_tracing.tracer as tr
 
-from hydragnn.utils.print_utils import log
-from hydragnn.utils import nsplit
+from hydragnn.utils.print.print_utils import log
+from hydragnn.utils.distributed import nsplit
 
 try:
     from hydragnn.utils.adiosdataset import AdiosDataset
@@ -151,7 +151,7 @@ def main():
 
     ##################################################################################################################
     # Always initialize for multi-rank training.
-    comm_size, rank = hydragnn.utils.setup_ddp()
+    comm_size, rank = hydragnn.utils.distributed.setup_ddp()
     ##################################################################################################################
 
     comm = MPI.COMM_WORLD
@@ -164,8 +164,8 @@ def main():
     )
 
     log_name = "gfm_test" if args.log is None else args.log
-    hydragnn.utils.setup_log(log_name)
-    writer = hydragnn.utils.get_summary_writer(log_name)
+    hydragnn.utils.print.setup_log(log_name)
+    writer = hydragnn.utils.model.get_summary_writer(log_name)
 
     log("Command: {0}\n".format(" ".join([x for x in sys.argv])), rank=0)
 
@@ -216,7 +216,7 @@ def main():
             trainset.pna_deg = pna_deg
     elif args.format == "multi":
         ## Reading multiple datasets, which requires the following arguments:
-        ## --multi_model_list: the list datasets/model names
+        ## --multi_model_list: the list dataset/model names
         modellist = args.multi_model_list.split(",")
         if rank == 0:
             ndata_list = list()
@@ -361,11 +361,13 @@ def main():
         trainset, valset, testset, config["NeuralNetwork"]["Training"]["batch_size"]
     )
 
-    config = hydragnn.utils.update_config(config, train_loader, val_loader, test_loader)
+    config = hydragnn.utils.input_config_parsing.update_config(
+        config, train_loader, val_loader, test_loader
+    )
     ## Good to sync with everyone right after DDStore setup
     comm.Barrier()
 
-    hydragnn.utils.save_config(config, log_name)
+    hydragnn.utils.input_config_parsing.save_config(config, log_name)
 
     timer.stop()
 
@@ -373,7 +375,7 @@ def main():
         config=config["NeuralNetwork"],
         verbosity=verbosity,
     )
-    model = hydragnn.utils.get_distributed_model(model, verbosity)
+    model = hydragnn.utils.distributed.get_distributed_model(model, verbosity)
 
     # Print details of neural network architecture
     print_model(model)
@@ -384,7 +386,7 @@ def main():
         optimizer, mode="min", factor=0.5, patience=5, min_lr=0.00001
     )
 
-    hydragnn.utils.load_existing_model_config(
+    hydragnn.utils.model.load_existing_model_config(
         model, config["NeuralNetwork"]["Training"], optimizer=optimizer
     )
 
@@ -404,8 +406,8 @@ def main():
         create_plots=False,
     )
 
-    hydragnn.utils.save_model(model, optimizer, log_name)
-    hydragnn.utils.print_timers(verbosity)
+    hydragnn.utils.model.save_model(model, optimizer, log_name)
+    hydragnn.utils.profiling_and_tracing.print_timers(verbosity)
 
     if tr.has("GPTLTracer"):
         import gptl4py as gp
diff --git a/examples/ogb/train_gap.py b/examples/ogb/train_gap.py
index 52174b26f..5b54763b3 100644
--- a/examples/ogb/train_gap.py
+++ b/examples/ogb/train_gap.py
@@ -8,30 +8,31 @@
 import sys
 from tqdm import tqdm
 from mpi4py import MPI
-from itertools import chain
 import argparse
-import time
 import math
 
 import hydragnn
 from hydragnn.preprocess.load_data import split_dataset
-from hydragnn.utils.print_utils import print_distributed, iterate_tqdm
-from hydragnn.utils.time_utils import Timer
-from hydragnn.utils.pickledataset import SimplePickleWriter, SimplePickleDataset
-from hydragnn.preprocess.utils import gather_deg
+from hydragnn.utils.print.print_utils import print_distributed
+from hydragnn.utils.profiling_and_tracing.time_utils import Timer
+from hydragnn.utils.datasets.pickledataset import (
+    SimplePickleWriter,
+    SimplePickleDataset,
+)
+from hydragnn.preprocess.graph_samples_checks_and_updates import gather_deg
 from hydragnn.utils.model import print_model
-from hydragnn.utils.smiles_utils import (
+from hydragnn.utils.descriptors_and_embeddings.smiles_utils import (
     get_node_attribute_name,
     generate_graphdata_from_smilestr,
 )
-from hydragnn.utils.config_utils import parse_deepspeed_config
+from hydragnn.utils.input_config_parsing.config_utils import parse_deepspeed_config
 from hydragnn.utils.distributed import get_deepspeed_init_args
-from hydragnn.utils import nsplit
+from hydragnn.utils.distributed import nsplit
 
 import numpy as np
 
 try:
-    from hydragnn.utils.adiosdataset import AdiosWriter, AdiosDataset
+    from hydragnn.utils.datasets.adiosdataset import AdiosWriter, AdiosDataset
 except ImportError:
     pass
 
@@ -92,7 +93,7 @@ def info(*args, logtype="info", sep=" "):
     getattr(logging, logtype)(sep.join(map(str, args)))
 
 
-from hydragnn.utils.abstractbasedataset import AbstractBaseDataset
+from hydragnn.utils.datasets import AbstractBaseDataset
 
 
 def smiles_to_graph(datadir, files_list):
@@ -136,7 +137,7 @@ def smiles_to_graph(datadir, files_list):
 
 
 class OGBDataset(AbstractBaseDataset):
-    """OGBDataset dataset class"""
+    """OGBDataset datasets class"""
 
     def __init__(self, dirpath, var_config, dist=False):
         super().__init__()
@@ -314,7 +315,9 @@ def __getitem__(self, idx):
     var_config["node_feature_dims"] = var_config["input_node_feature_dims"]
     ##################################################################################################################
     # Always initialize for multi-rank training.
-    comm_size, rank = hydragnn.utils.setup_ddp(use_deepspeed=args.use_deepspeed)
+    comm_size, rank = hydragnn.utils.distributed.setup_ddp(
+        use_deepspeed=args.use_deepspeed
+    )
     ##################################################################################################################
 
     comm = MPI.COMM_WORLD
@@ -327,9 +330,9 @@ def __getitem__(self, idx):
     )
 
     log_name = "ogb_" + inputfilesubstr
-    hydragnn.utils.setup_log(log_name)
-    writer = hydragnn.utils.get_summary_writer(log_name)
-    hydragnn.utils.save_config(config, log_name)
+    hydragnn.utils.print.setup_log(log_name)
+    writer = hydragnn.utils.model.get_summary_writer(log_name)
+    hydragnn.utils.input_config_parsing.save_config(config, log_name)
 
     modelname = "ogb_" + inputfilesubstr
     if args.preonly:
@@ -445,7 +448,9 @@ def __getitem__(self, idx):
         trainset, valset, testset, config["NeuralNetwork"]["Training"]["batch_size"]
     )
 
-    config = hydragnn.utils.update_config(config, train_loader, val_loader, test_loader)
+    config = hydragnn.utils.input_config_parsing.update_config(
+        config, train_loader, val_loader, test_loader
+    )
     timer.stop()
 
     model = hydragnn.models.create_model_config(
@@ -458,7 +463,7 @@ def __getitem__(self, idx):
     dist.barrier()
 
     if not args.use_deepspeed:
-        model = hydragnn.utils.get_distributed_model(model, verbosity)
+        model = hydragnn.utils.distributed.get_distributed_model(model, verbosity)
 
         learning_rate = config["NeuralNetwork"]["Training"]["Optimizer"][
             "learning_rate"
@@ -468,7 +473,7 @@ def __getitem__(self, idx):
             optimizer, mode="min", factor=0.5, patience=5, min_lr=0.00001
         )
 
-        hydragnn.utils.load_existing_model_config(
+        hydragnn.utils.model.load_existing_model_config(
             model, config["NeuralNetwork"]["Training"], optimizer=optimizer
         )
 
@@ -495,7 +500,7 @@ def __getitem__(self, idx):
             optimizer=optimizer,
         )  # scheduler is not managed by deepspeed because it is per-epoch instead of per-step
 
-        hydragnn.utils.load_existing_model_config(
+        hydragnn.utils.model.load_existing_model_config(
             model, config["NeuralNetwork"]["Training"], use_deepspeed=True
         )
 
@@ -516,8 +521,8 @@ def __getitem__(self, idx):
         use_deepspeed=args.use_deepspeed,
     )
 
-    hydragnn.utils.save_model(model, optimizer, log_name)
-    hydragnn.utils.print_timers(verbosity)
+    hydragnn.utils.model.save_model(model, optimizer, log_name)
+    hydragnn.utils.profiling_and_tracing.print_timers(verbosity)
 
     if args.mae:
         ##################################################################################################################
diff --git a/examples/open_catalyst_2020/download_dataset.py b/examples/open_catalyst_2020/download_dataset.py
index 99865f3ae..c938a07dc 100644
--- a/examples/open_catalyst_2020/download_dataset.py
+++ b/examples/open_catalyst_2020/download_dataset.py
@@ -140,8 +140,8 @@ def cleanup(filename, dirname):
     parser.add_argument(
         "--data-path",
         type=str,
-        default="./dataset",
-        help="Specify path to save dataset. Defaults to './dataset'",
+        default="./datasets",
+        help="Specify path to save datasets. Defaults to './datasets'",
     )
 
     args, _ = parser.parse_known_args()
diff --git a/examples/open_catalyst_2020/train.py b/examples/open_catalyst_2020/train.py
index d971fa15d..c920d9002 100644
--- a/examples/open_catalyst_2020/train.py
+++ b/examples/open_catalyst_2020/train.py
@@ -13,17 +13,20 @@
 torch.manual_seed(random_state)
 
 import hydragnn
-from hydragnn.utils.time_utils import Timer
+from hydragnn.utils.profiling_and_tracing.time_utils import Timer
 from hydragnn.utils.model import print_model
-from hydragnn.utils.abstractbasedataset import AbstractBaseDataset
-from hydragnn.utils.distdataset import DistDataset
-from hydragnn.utils.pickledataset import SimplePickleWriter, SimplePickleDataset
-from hydragnn.preprocess.utils import gather_deg
+from hydragnn.utils.datasets.abstractbasedataset import AbstractBaseDataset
+from hydragnn.utils.datasets.distdataset import DistDataset
+from hydragnn.utils.datasets.pickledataset import (
+    SimplePickleWriter,
+    SimplePickleDataset,
+)
+from hydragnn.preprocess.graph_samples_checks_and_updates import gather_deg
 from hydragnn.preprocess.load_data import split_dataset
 
-import hydragnn.utils.tracer as tr
+import hydragnn.utils.profiling_and_tracing.tracer as tr
 
-from hydragnn.utils.print_utils import iterate_tqdm, log
+from hydragnn.utils.print.print_utils import iterate_tqdm, log
 
 from utils.atoms_to_graphs import AtomsToGraphs
 from utils.preprocess import write_images_to_adios
@@ -34,7 +37,7 @@
     pass
 
 import subprocess
-from hydragnn.utils import nsplit
+from hydragnn.utils.distributed import nsplit
 
 ## FIMME
 torch.backends.cudnn.enabled = False
@@ -206,7 +209,7 @@ def get(self, idx):
 
     ##################################################################################################################
     # Always initialize for multi-rank training.
-    comm_size, rank = hydragnn.utils.setup_ddp()
+    comm_size, rank = hydragnn.utils.distributed.setup_ddp()
     ##################################################################################################################
 
     comm = MPI.COMM_WORLD
@@ -219,8 +222,8 @@ def get(self, idx):
     )
 
     log_name = "OC2020" if args.log is None else args.log
-    hydragnn.utils.setup_log(log_name)
-    writer = hydragnn.utils.get_summary_writer(log_name)
+    hydragnn.utils.print.setup_log(log_name)
+    writer = hydragnn.utils.model.get_summary_writer(log_name)
 
     log("Command: {0}\n".format(" ".join([x for x in sys.argv])), rank=0)
 
@@ -360,11 +363,13 @@ def get(self, idx):
         trainset, valset, testset, config["NeuralNetwork"]["Training"]["batch_size"]
     )
 
-    config = hydragnn.utils.update_config(config, train_loader, val_loader, test_loader)
+    config = hydragnn.utils.input_config_parsing.update_config(
+        config, train_loader, val_loader, test_loader
+    )
     ## Good to sync with everyone right after DDStore setup
     comm.Barrier()
 
-    hydragnn.utils.save_config(config, log_name)
+    hydragnn.utils.input_config_parsing.save_config(config, log_name)
 
     timer.stop()
 
@@ -372,7 +377,7 @@ def get(self, idx):
         config=config["NeuralNetwork"],
         verbosity=verbosity,
     )
-    model = hydragnn.utils.get_distributed_model(model, verbosity)
+    model = hydragnn.utils.distributed.get_distributed_model(model, verbosity)
 
     # Print details of neural network architecture
     print_model(model)
@@ -383,7 +388,7 @@ def get(self, idx):
         optimizer, mode="min", factor=0.5, patience=5, min_lr=0.00001
     )
 
-    hydragnn.utils.load_existing_model_config(
+    hydragnn.utils.model.load_existing_model_config(
         model, config["NeuralNetwork"]["Training"], optimizer=optimizer
     )
 
@@ -403,8 +408,8 @@ def get(self, idx):
         create_plots=False,
     )
 
-    hydragnn.utils.save_model(model, optimizer, log_name)
-    hydragnn.utils.print_timers(verbosity)
+    hydragnn.utils.model.save_model(model, optimizer, log_name)
+    hydragnn.utils.profiling_and_tracing.print_timers(verbosity)
 
     if tr.has("GPTLTracer"):
         import gptl4py as gp
diff --git a/examples/open_catalyst_2020/uncompress.py b/examples/open_catalyst_2020/uncompress.py
index 49f223c81..8f7a92328 100644
--- a/examples/open_catalyst_2020/uncompress.py
+++ b/examples/open_catalyst_2020/uncompress.py
@@ -28,7 +28,7 @@ def decompress_list_of_files(ip_op_pair: Tuple[str, str]) -> None:
 def get_parser() -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser()
     parser.add_argument(
-        "--ipdir", type=str, help="Path to compressed dataset directory"
+        "--ipdir", type=str, help="Path to compressed datasets directory"
     )
     parser.add_argument(
         "--opdir", type=str, help="Directory path to uncompress files to"
diff --git a/examples/open_catalyst_2020/utils/atoms_to_graphs.py b/examples/open_catalyst_2020/utils/atoms_to_graphs.py
index 10222757c..a0edc97ce 100644
--- a/examples/open_catalyst_2020/utils/atoms_to_graphs.py
+++ b/examples/open_catalyst_2020/utils/atoms_to_graphs.py
@@ -15,7 +15,10 @@
 from torch_geometric.data import Data
 from torch_geometric.transforms import Distance, Spherical, LocalCartesian
 
-from hydragnn.preprocess.utils import RadiusGraph, RadiusGraphPBC
+from hydragnn.preprocess.graph_samples_checks_and_updates import (
+    RadiusGraph,
+    RadiusGraphPBC,
+)
 
 # transform_coordinates = Spherical(norm=False, cat=False)
 # transform_coordinates = LocalCartesian(norm=False, cat=False)
diff --git a/examples/open_catalyst_2022/train.py b/examples/open_catalyst_2022/train.py
index 80d098852..ea109222a 100644
--- a/examples/open_catalyst_2022/train.py
+++ b/examples/open_catalyst_2022/train.py
@@ -19,19 +19,25 @@
 from torch_geometric.transforms import Distance, Spherical, LocalCartesian
 
 import hydragnn
-from hydragnn.utils.time_utils import Timer
+from hydragnn.utils.profiling_and_tracing.time_utils import Timer
 from hydragnn.utils.model import print_model
-from hydragnn.utils.abstractbasedataset import AbstractBaseDataset
-from hydragnn.utils.distdataset import DistDataset
-from hydragnn.utils.pickledataset import SimplePickleWriter, SimplePickleDataset
-from hydragnn.preprocess.utils import gather_deg
+from hydragnn.utils.datasets.abstractbasedataset import AbstractBaseDataset
+from hydragnn.utils.datasets.distdataset import DistDataset
+from hydragnn.utils.datasets.pickledataset import (
+    SimplePickleWriter,
+    SimplePickleDataset,
+)
+from hydragnn.preprocess.graph_samples_checks_and_updates import gather_deg
 from hydragnn.preprocess.load_data import split_dataset
 
-import hydragnn.utils.tracer as tr
+import hydragnn.utils.profiling_and_tracing.tracer as tr
 
-from hydragnn.utils.print_utils import iterate_tqdm, log
+from hydragnn.utils.print.print_utils import iterate_tqdm, log
 
-from hydragnn.preprocess.utils import RadiusGraph, RadiusGraphPBC
+from hydragnn.preprocess.graph_samples_checks_and_updates import (
+    RadiusGraph,
+    RadiusGraphPBC,
+)
 
 from ase.io import read
 
@@ -40,8 +46,7 @@
 except ImportError:
     pass
 
-import subprocess
-from hydragnn.utils import nsplit
+from hydragnn.utils.distributed import nsplit
 
 
 def info(*args, logtype="info", sep=" "):
@@ -259,7 +264,7 @@ def get(self, idx):
 
     ##################################################################################################################
     # Always initialize for multi-rank training.
-    comm_size, rank = hydragnn.utils.setup_ddp()
+    comm_size, rank = hydragnn.utils.distributed.setup_ddp()
     ##################################################################################################################
 
     comm = MPI.COMM_WORLD
@@ -272,8 +277,8 @@ def get(self, idx):
     )
 
     log_name = "OC2022" if args.log is None else args.log
-    hydragnn.utils.setup_log(log_name)
-    writer = hydragnn.utils.get_summary_writer(log_name)
+    hydragnn.utils.print.setup_log(log_name)
+    writer = hydragnn.utils.model.get_summary_writer(log_name)
 
     log("Command: {0}\n".format(" ".join([x for x in sys.argv])), rank=0)
 
@@ -375,7 +380,7 @@ def get(self, idx):
             "ddstore": args.ddstore,
             "ddstore_width": args.ddstore_width,
         }
-        fname = os.path.join(os.path.dirname(__file__), "./dataset/%s.bp" % modelname)
+        fname = os.path.join(os.path.dirname(__file__), ".//%s.bp" % modelname)
         trainset = AdiosDataset(fname, "trainset", comm, **opt, var_config=var_config)
         valset = AdiosDataset(fname, "valset", comm, **opt, var_config=var_config)
         testset = AdiosDataset(fname, "testset", comm, **opt, var_config=var_config)
@@ -420,11 +425,13 @@ def get(self, idx):
         trainset, valset, testset, config["NeuralNetwork"]["Training"]["batch_size"]
     )
 
-    config = hydragnn.utils.update_config(config, train_loader, val_loader, test_loader)
+    config = hydragnn.utils.input_config_parsing.update_config(
+        config, train_loader, val_loader, test_loader
+    )
     ## Good to sync with everyone right after DDStore setup
     comm.Barrier()
 
-    hydragnn.utils.save_config(config, log_name)
+    hydragnn.utils.input_config_parsing.save_config(config, log_name)
 
     timer.stop()
 
@@ -432,7 +439,7 @@ def get(self, idx):
         config=config["NeuralNetwork"],
         verbosity=verbosity,
     )
-    model = hydragnn.utils.get_distributed_model(model, verbosity)
+    model = hydragnn.utils.distributed.get_distributed_model(model, verbosity)
 
     # Print details of neural network architecture
     print_model(model)
@@ -443,7 +450,7 @@ def get(self, idx):
         optimizer, mode="min", factor=0.5, patience=5, min_lr=0.00001
     )
 
-    hydragnn.utils.load_existing_model_config(
+    hydragnn.utils.model.load_existing_model_config(
         model, config["NeuralNetwork"]["Training"], optimizer=optimizer
     )
 
@@ -463,8 +470,8 @@ def get(self, idx):
         create_plots=False,
     )
 
-    hydragnn.utils.save_model(model, optimizer, log_name)
-    hydragnn.utils.print_timers(verbosity)
+    hydragnn.utils.model.save_model(model, optimizer, log_name)
+    hydragnn.utils.profiling_and_tracing.print_timers(verbosity)
 
     if tr.has("GPTLTracer"):
         import gptl4py as gp
diff --git a/examples/qm7x/train.py b/examples/qm7x/train.py
index 12a176da4..e8e9f0334 100644
--- a/examples/qm7x/train.py
+++ b/examples/qm7x/train.py
@@ -14,14 +14,17 @@
 import argparse
 
 import hydragnn
-from hydragnn.utils.print_utils import iterate_tqdm, log
-from hydragnn.utils.time_utils import Timer
+from hydragnn.utils.print.print_utils import iterate_tqdm, log
+from hydragnn.utils.profiling_and_tracing.time_utils import Timer
 
 from hydragnn.utils.distributed import get_device
 from hydragnn.preprocess.load_data import split_dataset
-from hydragnn.utils.distdataset import DistDataset
-from hydragnn.utils.pickledataset import SimplePickleWriter, SimplePickleDataset
-from hydragnn.preprocess.utils import gather_deg
+from hydragnn.utils.datasets.distdataset import DistDataset
+from hydragnn.utils.datasets.pickledataset import (
+    SimplePickleWriter,
+    SimplePickleDataset,
+)
+from hydragnn.preprocess.graph_samples_checks_and_updates import gather_deg
 
 import numpy as np
 
@@ -42,8 +45,8 @@
 except ImportError:
     pass
 
-from hydragnn.utils import nsplit
-import hydragnn.utils.tracer as tr
+from hydragnn.utils.distributed import nsplit
+import hydragnn.utils.profiling_and_tracing.tracer as tr
 
 # FIXME: this works fine for now because we train on QM7-X molecules
 # for larger chemical spaces, the following atom representation has to be properly expanded
@@ -65,7 +68,7 @@ def info(*args, logtype="info", sep=" "):
     getattr(logging, logtype)(sep.join(map(str, args)))
 
 
-from hydragnn.utils.abstractbasedataset import AbstractBaseDataset
+from hydragnn.utils.datasets.abstractbasedataset import AbstractBaseDataset
 
 # FIXME: this radis cutoff overwrites the radius cutoff currently written in the JSON file
 create_graph_fromXYZ = RadiusGraph(r=5.0)  # radius cutoff in angstrom
@@ -73,7 +76,7 @@ def info(*args, logtype="info", sep=" "):
 
 
 class QM7XDataset(AbstractBaseDataset):
-    """QM7-XDataset dataset class"""
+    """QM7-XDataset datasets class"""
 
     def __init__(self, dirpath, var_config, energy_per_atom=True, dist=False):
         super().__init__()
@@ -302,7 +305,7 @@ def get(self, idx):
 
     ##################################################################################################################
     # Always initialize for multi-rank training.
-    comm_size, rank = hydragnn.utils.setup_ddp()
+    comm_size, rank = hydragnn.utils.distributed.setup_ddp()
     ##################################################################################################################
 
     comm = MPI.COMM_WORLD
@@ -315,8 +318,8 @@ def get(self, idx):
     )
 
     log_name = "qm7x" if args.log is None else args.log
-    hydragnn.utils.setup_log(log_name)
-    writer = hydragnn.utils.get_summary_writer(log_name)
+    hydragnn.utils.print.setup_log(log_name)
+    writer = hydragnn.utils.model.get_summary_writer(log_name)
 
     log("Command: {0}\n".format(" ".join([x for x in sys.argv])), rank=0)
 
@@ -451,11 +454,13 @@ def get(self, idx):
         trainset, valset, testset, config["NeuralNetwork"]["Training"]["batch_size"]
     )
 
-    config = hydragnn.utils.update_config(config, train_loader, val_loader, test_loader)
+    config = hydragnn.utils.input_config_parsing.update_config(
+        config, train_loader, val_loader, test_loader
+    )
     ## Good to sync with everyone right after DDStore setup
     comm.Barrier()
 
-    hydragnn.utils.save_config(config, log_name)
+    hydragnn.utils.input_config_parsing.save_config(config, log_name)
 
     timer.stop()
 
@@ -463,7 +468,7 @@ def get(self, idx):
         config=config["NeuralNetwork"],
         verbosity=verbosity,
     )
-    model = hydragnn.utils.get_distributed_model(model, verbosity)
+    model = hydragnn.utils.distributed.get_distributed_model(model, verbosity)
 
     learning_rate = config["NeuralNetwork"]["Training"]["Optimizer"]["learning_rate"]
     optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
@@ -471,7 +476,7 @@ def get(self, idx):
         optimizer, mode="min", factor=0.5, patience=5, min_lr=0.00001
     )
 
-    hydragnn.utils.load_existing_model_config(
+    hydragnn.utils.model.load_existing_model_config(
         model, config["NeuralNetwork"]["Training"], optimizer=optimizer
     )
 
@@ -491,8 +496,8 @@ def get(self, idx):
         create_plots=False,
     )
 
-    hydragnn.utils.save_model(model, optimizer, log_name)
-    hydragnn.utils.print_timers(verbosity)
+    hydragnn.utils.model.save_model(model, optimizer, log_name)
+    hydragnn.utils.profiling_and_tracing.print_timers(verbosity)
 
     if tr.has("GPTLTracer"):
         import gptl4py as gp
diff --git a/examples/qm9/qm9.py b/examples/qm9/qm9.py
index 2e6c3da8e..1be651f97 100644
--- a/examples/qm9/qm9.py
+++ b/examples/qm9/qm9.py
@@ -47,13 +47,13 @@ def qm9_pre_filter(data):
 var_config = config["NeuralNetwork"]["Variables_of_interest"]
 
 # Always initialize for multi-rank training.
-world_size, world_rank = hydragnn.utils.setup_ddp()
+world_size, world_rank = hydragnn.utils.distributed.setup_ddp()
 
 log_name = "qm9_test"
 # Enable print to log file.
-hydragnn.utils.setup_log(log_name)
+hydragnn.utils.print.print_utils.setup_log(log_name)
 
-# Use built-in torch_geometric dataset.
+# Use built-in torch_geometric datasets.
 # Filter function above used to run quick example.
 # NOTE: data is moved to the device in the pre-transform.
 # NOTE: transforms/filters will NOT be re-run unless the qm9/processed/ directory is removed.
@@ -67,13 +67,15 @@ def qm9_pre_filter(data):
     train, val, test, config["NeuralNetwork"]["Training"]["batch_size"]
 )
 
-config = hydragnn.utils.update_config(config, train_loader, val_loader, test_loader)
+config = hydragnn.utils.input_config_parsing.update_config(
+    config, train_loader, val_loader, test_loader
+)
 
 model = hydragnn.models.create_model_config(
     config=config["NeuralNetwork"],
     verbosity=verbosity,
 )
-model = hydragnn.utils.get_distributed_model(model, verbosity)
+model = hydragnn.utils.distributed.get_distributed_model(model, verbosity)
 
 learning_rate = config["NeuralNetwork"]["Training"]["Optimizer"]["learning_rate"]
 optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
@@ -81,9 +83,9 @@ def qm9_pre_filter(data):
     optimizer, mode="min", factor=0.5, patience=5, min_lr=0.00001
 )
 
-# Run training with the given model and qm9 dataset.
-writer = hydragnn.utils.get_summary_writer(log_name)
-hydragnn.utils.save_config(config, log_name)
+# Run training with the given model and qm9 datasets.
+writer = hydragnn.utils.model.model.get_summary_writer(log_name)
+hydragnn.utils.input_config_parsing.save_config(config, log_name)
 
 hydragnn.train.train_validate_test(
     model,
diff --git a/examples/qm9_hpo/qm9.py b/examples/qm9_hpo/qm9.py
index 61ef2376e..95ace3382 100644
--- a/examples/qm9_hpo/qm9.py
+++ b/examples/qm9_hpo/qm9.py
@@ -75,7 +75,7 @@ def qm9_pre_filter(data):
 # Enable print to log file.
 hydragnn.utils.setup_log(log_name)
 
-# Use built-in torch_geometric dataset.
+# Use built-in torch_geometric datasets.
 # Filter function above used to run quick example.
 # NOTE: data is moved to the device in the pre-transform.
 # NOTE: transforms/filters will NOT be re-run unless the qm9/processed/ directory is removed.
@@ -103,7 +103,7 @@ def qm9_pre_filter(data):
     optimizer, mode="min", factor=0.5, patience=5, min_lr=0.00001
 )
 
-# Run training with the given model and qm9 dataset.
+# Run training with the given model and qm9 datasets.
 writer = hydragnn.utils.get_summary_writer(log_name)
 hydragnn.utils.save_config(config, log_name)
 
diff --git a/examples/qm9_hpo/qm9_deephyper.py b/examples/qm9_hpo/qm9_deephyper.py
index cb4019cf4..51031c27a 100644
--- a/examples/qm9_hpo/qm9_deephyper.py
+++ b/examples/qm9_hpo/qm9_deephyper.py
@@ -135,7 +135,7 @@ def run(trial):
 
     log_name = "qm9"
 
-    # Use built-in torch_geometric dataset.
+    # Use built-in torch_geometric datasets.
     # Filter function above used to run quick example.
     # NOTE: data is moved to the device in the pre-transform.
     # NOTE: transforms/filters will NOT be re-run unless the qm9/processed/ directory is removed.
diff --git a/examples/qm9_hpo/qm9_optuna.py b/examples/qm9_hpo/qm9_optuna.py
index 64403ac29..0b57f9e83 100644
--- a/examples/qm9_hpo/qm9_optuna.py
+++ b/examples/qm9_hpo/qm9_optuna.py
@@ -171,7 +171,7 @@ def objective(trial):
     # Enable print to log file.
     hydragnn.utils.setup_log(log_name)
 
-    # Use built-in torch_geometric dataset.
+    # Use built-in torch_geometric datasets.
     # Filter function above used to run quick example.
     # NOTE: data is moved to the device in the pre-transform.
     # NOTE: transforms/filters will NOT be re-run unless the qm9/processed/ directory is removed.
diff --git a/hydragnn/models/Base.py b/hydragnn/models/Base.py
index 2bcb791ba..186d0209f 100644
--- a/hydragnn/models/Base.py
+++ b/hydragnn/models/Base.py
@@ -19,7 +19,7 @@
 from hydragnn.utils.model import activation_function_selection, loss_function_selection
 import sys
 from hydragnn.utils.distributed import get_device
-from hydragnn.utils.print_utils import print_master
+from hydragnn.utils.print.print_utils import print_master
 
 import inspect
 
@@ -264,7 +264,7 @@ def _multihead(self):
                     assert (
                         self.num_nodes is not None
                     ), "num_nodes must be positive integer for MLP"
-                    # """if different graphs in the dataset have different size, one MLP is shared across all nodes """
+                    # """if different graphs in the datasets have different size, one MLP is shared across all nodes """
                     head_NN = MLPNode(
                         self.hidden_dim,
                         self.head_dims[ihead] * (1 + self.var_output),
diff --git a/hydragnn/models/EGCLStack.py b/hydragnn/models/EGCLStack.py
index 639196555..7109d0fc3 100644
--- a/hydragnn/models/EGCLStack.py
+++ b/hydragnn/models/EGCLStack.py
@@ -15,7 +15,7 @@
 from torch_geometric.nn import Sequential
 from .Base import Base
 
-from ..utils import unsorted_segment_mean
+from hydragnn.utils.model import unsorted_segment_mean
 
 
 class EGCLStack(Base):
diff --git a/hydragnn/models/SCFStack.py b/hydragnn/models/SCFStack.py
index 7c67cffd5..4f66ae6b5 100644
--- a/hydragnn/models/SCFStack.py
+++ b/hydragnn/models/SCFStack.py
@@ -26,7 +26,7 @@
 
 from .Base import Base
 
-from ..utils import unsorted_segment_mean
+from hydragnn.utils.model import unsorted_segment_mean
 
 
 class SCFStack(Base):
diff --git a/hydragnn/models/create.py b/hydragnn/models/create.py
index 4e4f6e2f3..c39f1414f 100644
--- a/hydragnn/models/create.py
+++ b/hydragnn/models/create.py
@@ -25,8 +25,7 @@
 from hydragnn.models.EGCLStack import EGCLStack
 
 from hydragnn.utils.distributed import get_device
-from hydragnn.utils.print_utils import print_distributed
-from hydragnn.utils.time_utils import Timer
+from hydragnn.utils.profiling_and_tracing.time_utils import Timer
 
 
 def create_model_config(
diff --git a/hydragnn/preprocess/__init__.py b/hydragnn/preprocess/__init__.py
index c5ce6ac20..ccc0d2bb0 100644
--- a/hydragnn/preprocess/__init__.py
+++ b/hydragnn/preprocess/__init__.py
@@ -1,6 +1,6 @@
 from .dataset_descriptors import AtomFeatures, StructureFeatures
 
-from .utils import (
+from .graph_samples_checks_and_updates import (
     check_if_graph_size_variable,
     check_if_graph_size_variable_dist,
     get_radius_graph,
@@ -10,9 +10,10 @@
     RadiusGraphPBC,
     update_predicted_values,
     update_atom_features,
-    stratified_sampling,
 )
 
+from .stratified_sampling import stratified_sampling
+
 from .load_data import (
     dataset_loading_and_splitting,
     create_dataloaders,
@@ -26,4 +27,3 @@
 )
 from .lsms_raw_dataset_loader import LSMS_RawDataLoader
 from .cfg_raw_dataset_loader import CFG_RawDataLoader
-from .compositional_data_splitting import compositional_stratified_splitting
diff --git a/hydragnn/preprocess/cfg_raw_dataset_loader.py b/hydragnn/preprocess/cfg_raw_dataset_loader.py
index b5043abb1..32f44c867 100644
--- a/hydragnn/preprocess/cfg_raw_dataset_loader.py
+++ b/hydragnn/preprocess/cfg_raw_dataset_loader.py
@@ -12,7 +12,6 @@
 import os
 import numpy as np
 
-import torch
 from torch_geometric.data import Data
 from torch import tensor
 
diff --git a/hydragnn/preprocess/utils.py b/hydragnn/preprocess/graph_samples_checks_and_updates.py
similarity index 85%
rename from hydragnn/preprocess/utils.py
rename to hydragnn/preprocess/graph_samples_checks_and_updates.py
index 3533756d2..b4162d742 100644
--- a/hydragnn/preprocess/utils.py
+++ b/hydragnn/preprocess/graph_samples_checks_and_updates.py
@@ -20,7 +20,7 @@
 
 from .dataset_descriptors import AtomFeatures
 
-## This function can be slow if dataset is too large. Use with caution.
+## This function can be slow if datasets is too large. Use with caution.
 ## Recommend to use check_if_graph_size_variable_dist
 def check_if_graph_size_variable(train_loader, val_loader, test_loader):
     backend = os.getenv("HYDRAGNN_AGGR_BACKEND", "torch")
@@ -175,7 +175,7 @@ def __repr__(self) -> str:
 
 
 def gather_deg(dataset):
-    from hydragnn.utils.print_utils import iterate_tqdm
+    from hydragnn.utils.print.print_utils import iterate_tqdm
 
     backend = os.getenv("HYDRAGNN_AGGR_BACKEND", "torch")
     if backend == "torch":
@@ -197,7 +197,7 @@ def gather_deg(dataset):
 
 def gather_deg_dist(dataset):
     import torch.distributed as dist
-    from hydragnn.utils.print_utils import iterate_tqdm
+    from hydragnn.utils.print.print_utils import iterate_tqdm
     from hydragnn.utils.distributed import get_device
 
     max_deg = 0
@@ -218,7 +218,7 @@ def gather_deg_dist(dataset):
 
 def gather_deg_mpi(dataset):
     from mpi4py import MPI
-    from hydragnn.utils.print_utils import iterate_tqdm
+    from hydragnn.utils.print.print_utils import iterate_tqdm
 
     max_deg = 0
     for data in iterate_tqdm(dataset, 2, desc="Degree max"):
@@ -290,47 +290,3 @@ def update_atom_features(atom_features: [AtomFeatures], data: Data):
     """
     feature_indices = [i for i in atom_features]
     data.x = data.x[:, feature_indices]
-
-
-def stratified_sampling(dataset: [Data], subsample_percentage: float, verbosity=0):
-    """Given the dataset and the percentage of data you want to extract from it, method will
-    apply stratified sampling where X is the dataset and Y is are the category values for each datapoint.
-    In the case of the structures dataset where each structure contains 2 types of atoms, the category will
-    be constructed in a way: number of atoms of type 1 + number of protons of type 2 * 100.
-
-    Parameters
-    ----------
-    dataset: [Data]
-        A list of Data objects representing a structure that has atoms.
-    subsample_percentage: float
-        Percentage of the dataset.
-
-    Returns
-    ----------
-    [Data]
-        Subsample of the original dataset constructed using stratified sampling.
-    """
-    dataset_categories = []
-    print_distributed(verbosity, "Computing the categories for the whole dataset.")
-    for data in iterate_tqdm(dataset, verbosity):
-        frequencies = torch.bincount(data.x[:, 0].int())
-        frequencies = sorted(frequencies[frequencies > 0].tolist())
-        category = 0
-        for index, frequency in enumerate(frequencies):
-            category += frequency * (100 ** index)
-        dataset_categories.append(category)
-
-    subsample_indices = []
-    subsample = []
-
-    sss = StratifiedShuffleSplit(
-        n_splits=1, train_size=subsample_percentage, random_state=0
-    )
-
-    for subsample_index, rest_of_data_index in sss.split(dataset, dataset_categories):
-        subsample_indices = subsample_index.tolist()
-
-    for index in subsample_indices:
-        subsample.append(dataset[index])
-
-    return subsample
diff --git a/hydragnn/preprocess/load_data.py b/hydragnn/preprocess/load_data.py
index 606c4a680..90ef91f87 100644
--- a/hydragnn/preprocess/load_data.py
+++ b/hydragnn/preprocess/load_data.py
@@ -16,7 +16,6 @@
 
 import torch
 import torch.distributed as dist
-import torch_geometric
 
 # FIXME: deprecated in torch_geometric 2.0
 try:
@@ -27,23 +26,21 @@
 from hydragnn.preprocess.serialized_dataset_loader import SerializedDataLoader
 from hydragnn.preprocess.lsms_raw_dataset_loader import LSMS_RawDataLoader
 from hydragnn.preprocess.cfg_raw_dataset_loader import CFG_RawDataLoader
-from hydragnn.preprocess.compositional_data_splitting import (
+from hydragnn.utils.datasets.compositional_data_splitting import (
     compositional_stratified_splitting,
 )
 from hydragnn.utils.distributed import get_comm_size_and_rank
-from hydragnn.utils.time_utils import Timer
+from hydragnn.utils.profiling_and_tracing.time_utils import Timer
 import pickle
 
-from hydragnn.utils.print_utils import print_master, log
+from hydragnn.utils.print.print_utils import log
 
-from torch_geometric.data import Batch, Dataset
+from torch_geometric.data import Batch
 from torch.utils.data.dataloader import _DatasetKind
 
-from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
+from concurrent.futures import ThreadPoolExecutor
 import multiprocessing as mp
 import queue
-import time
-import sys
 import re
 
 
@@ -211,7 +208,7 @@ def dataset_loading_and_splitting(config: {}):
     if not list(config["Dataset"]["path"].values())[0].endswith(".pkl"):
         transform_raw_data_to_serialized(config["Dataset"])
 
-    ##if total dataset is provided, split the dataset and save them to pkl files and update config with pkl file locations
+    ##if total datasets is provided, split the datasets and save them to pkl files and update config with pkl file locations
     if "total" in config["Dataset"]["path"].keys():
         total_to_train_val_test_pkls(config)
 
@@ -374,7 +371,7 @@ def total_to_train_val_test_pkls(config, isdist=False):
         file_dir = config["Dataset"]["path"]["total"]
     else:
         file_dir = f"{os.environ['SERIALIZED_DATA_PATH']}/serialized_dataset/{config['Dataset']['name']}.pkl"
-    # if "total" raw dataset is provided, generate train/val/test pkl files and update config dict.
+    # if "total" raw datasets is provided, generate train/val/test pkl files and update config dict.
     with open(file_dir, "rb") as f:
         minmax_node_feature = pickle.load(f)
         minmax_graph_feature = pickle.load(f)
diff --git a/hydragnn/preprocess/raw_dataset_loader.py b/hydragnn/preprocess/raw_dataset_loader.py
index c0443bf2a..702e0ef92 100644
--- a/hydragnn/preprocess/raw_dataset_loader.py
+++ b/hydragnn/preprocess/raw_dataset_loader.py
@@ -14,12 +14,11 @@
 import pickle
 
 import torch
-from torch_geometric.data import Data
-from torch import tensor
 
 from hydragnn.utils.distributed import get_device
-from hydragnn.utils.print_utils import log
-from hydragnn.utils import nsplit, tensor_divide, comm_reduce
+from hydragnn.utils.print.print_utils import log
+from hydragnn.utils.distributed import nsplit, comm_reduce
+from hydragnn.utils.model.model import tensor_divide
 
 import random
 
@@ -38,7 +37,7 @@ class AbstractRawDataLoader:
     def __init__(self, config, dist=False):
         """
         config:
-          shows the dataset path the target variables information, e.g, location and dimension, in data file
+          shows the datasets path the target variables information, e.g, location and dimension, in data file
         ###########
         dataset_list:
           list of datasets read from self.path_dictionary
@@ -193,7 +192,7 @@ def scale_features_by_num_nodes(self, dataset):
 
     def normalize_dataset(self):
 
-        """Performs the normalization on Data objects and returns the normalized dataset."""
+        """Performs the normalization on Data objects and returns the normalized datasets."""
         num_node_features = len(self.node_feature_dim)
         num_graph_features = len(self.graph_feature_dim)
 
diff --git a/hydragnn/preprocess/serialized_dataset_loader.py b/hydragnn/preprocess/serialized_dataset_loader.py
index bef054edf..3b385f936 100644
--- a/hydragnn/preprocess/serialized_dataset_loader.py
+++ b/hydragnn/preprocess/serialized_dataset_loader.py
@@ -23,8 +23,8 @@
 
 from hydragnn.preprocess import update_predicted_values, update_atom_features
 from hydragnn.utils.distributed import get_device
-from hydragnn.utils.print_utils import print_distributed, iterate_tqdm
-from hydragnn.preprocess.utils import (
+from hydragnn.utils.print.print_utils import print_distributed, iterate_tqdm
+from hydragnn.preprocess.graph_samples_checks_and_updates import (
     get_radius_graph,
     get_radius_graph_pbc,
 )
@@ -194,9 +194,9 @@ def load_serialized_data(self, dataset_path: str):
         return dataset
 
     def __stratified_sampling(self, dataset: [Data], subsample_percentage: float):
-        """Given the dataset and the percentage of data you want to extract from it, method will
-        apply stratified sampling where X is the dataset and Y is are the category values for each datapoint.
-        In the case of the structures dataset where each structure contains 2 types of atoms, the category will
+        """Given the datasets and the percentage of data you want to extract from it, method will
+        apply stratified sampling where X is the datasets and Y is are the category values for each datapoint.
+        In the case of the structures datasets where each structure contains 2 types of atoms, the category will
         be constructed in a way: number of atoms of type 1 + number of protons of type 2 * 100.
 
         Parameters
@@ -204,16 +204,16 @@ def __stratified_sampling(self, dataset: [Data], subsample_percentage: float):
         dataset: [Data]
             A list of Data objects representing a structure that has atoms.
         subsample_percentage: float
-            Percentage of the dataset.
+            Percentage of the datasets.
 
         Returns
         ----------
         [Data]
-            Subsample of the original dataset constructed using stratified sampling.
+            Subsample of the original datasets constructed using stratified sampling.
         """
         dataset_categories = []
         print_distributed(
-            self.verbosity, "Computing the categories for the whole dataset."
+            self.verbosity, "Computing the categories for the whole datasets."
         )
         for data in iterate_tqdm(dataset, self.verbosity):
             frequencies = torch.bincount(data.x[:, 0].int())
diff --git a/hydragnn/preprocess/stratified_sampling.py b/hydragnn/preprocess/stratified_sampling.py
new file mode 100644
index 000000000..ebc68b2ed
--- /dev/null
+++ b/hydragnn/preprocess/stratified_sampling.py
@@ -0,0 +1,48 @@
+import torch
+from torch_geometric.data import Data
+from hydragnn.utils.print.print_utils import print_distributed, iterate_tqdm
+from sklearn.model_selection import StratifiedShuffleSplit
+
+
+def stratified_sampling(dataset: [Data], subsample_percentage: float, verbosity=0):
+    """Given the datasets and the percentage of data you want to extract from it, method will
+    apply stratified sampling where X is the datasets and Y is are the category values for each datapoint.
+    In the case of the structures datasets where each structure contains 2 types of atoms, the category will
+    be constructed in a way: number of atoms of type 1 + number of protons of type 2 * 100.
+
+    Parameters
+    ----------
+    dataset: [Data]
+        A list of Data objects representing a structure that has atoms.
+    subsample_percentage: float
+        Percentage of the datasets.
+
+    Returns
+    ----------
+    [Data]
+        Subsample of the original datasets constructed using stratified sampling.
+    """
+    dataset_categories = []
+    print_distributed(verbosity, "Computing the categories for the whole datasets.")
+    for data in iterate_tqdm(dataset, verbosity):
+        frequencies = torch.bincount(data.x[:, 0].int())
+        frequencies = sorted(frequencies[frequencies > 0].tolist())
+        category = 0
+        for index, frequency in enumerate(frequencies):
+            category += frequency * (100 ** index)
+        dataset_categories.append(category)
+
+    subsample_indices = []
+    subsample = []
+
+    sss = StratifiedShuffleSplit(
+        n_splits=1, train_size=subsample_percentage, random_state=0
+    )
+
+    for subsample_index, _ in sss.split(dataset, dataset_categories):
+        subsample_indices = subsample_index.tolist()
+
+    for index in subsample_indices:
+        subsample.append(dataset[index])
+
+    return subsample
diff --git a/hydragnn/run_prediction.py b/hydragnn/run_prediction.py
index b1b7bf3de..3d15f26a3 100755
--- a/hydragnn/run_prediction.py
+++ b/hydragnn/run_prediction.py
@@ -15,7 +15,7 @@
 from hydragnn.preprocess.load_data import dataset_loading_and_splitting
 from hydragnn.utils.distributed import setup_ddp, get_distributed_model
 from hydragnn.utils.model import load_existing_model
-from hydragnn.utils.config_utils import (
+from hydragnn.utils.input_config_parsing.config_utils import (
     update_config,
     get_log_name_config,
     parse_deepspeed_config,
diff --git a/hydragnn/run_training.py b/hydragnn/run_training.py
index 035693380..4c17f3e13 100644
--- a/hydragnn/run_training.py
+++ b/hydragnn/run_training.py
@@ -12,7 +12,6 @@
 import os, json
 from functools import singledispatch
 
-import torch
 import torch.distributed as dist
 from torch.optim.lr_scheduler import ReduceLROnPlateau
 
@@ -20,16 +19,16 @@
 from hydragnn.utils.distributed import (
     setup_ddp,
     get_distributed_model,
-    print_peak_memory,
 )
+from hydragnn.utils.distributed import print_peak_memory
 from hydragnn.utils.model import (
     save_model,
     get_summary_writer,
     load_existing_model_config,
 )
-from hydragnn.utils.print_utils import print_distributed, setup_log
-from hydragnn.utils.time_utils import print_timers
-from hydragnn.utils.config_utils import (
+from hydragnn.utils.print.print_utils import print_distributed, setup_log
+from hydragnn.utils.profiling_and_tracing.time_utils import print_timers
+from hydragnn.utils.input_config_parsing.config_utils import (
     update_config,
     get_log_name_config,
     save_config,
diff --git a/hydragnn/train/train_validate_test.py b/hydragnn/train/train_validate_test.py
index 3232da2a8..5195d5134 100644
--- a/hydragnn/train/train_validate_test.py
+++ b/hydragnn/train/train_validate_test.py
@@ -17,23 +17,21 @@
 from hydragnn.preprocess.serialized_dataset_loader import SerializedDataLoader
 from hydragnn.postprocess.postprocess import output_denormalize
 from hydragnn.postprocess.visualizer import Visualizer
-from hydragnn.utils.print_utils import print_distributed, iterate_tqdm, log
-from hydragnn.utils.time_utils import Timer
-from hydragnn.utils.profile import Profiler
-from hydragnn.utils.distributed import get_device, print_peak_memory, check_remaining
-from hydragnn.preprocess.load_data import HydraDataLoader
-from hydragnn.utils.model import Checkpoint, EarlyStopping
+from hydragnn.utils.print.print_utils import print_distributed, iterate_tqdm
+from hydragnn.utils.profiling_and_tracing.time_utils import Timer
+from hydragnn.utils.profiling_and_tracing.profile import Profiler
+from hydragnn.utils.distributed import get_device, check_remaining
+from hydragnn.utils.model.model import Checkpoint, EarlyStopping
 
 import os
 
 from torch.profiler import record_function
-import contextlib
-from unittest.mock import MagicMock
+
 from hydragnn.utils.distributed import get_comm_size_and_rank
 import torch.distributed as dist
 import pickle
 
-import hydragnn.utils.tracer as tr
+import hydragnn.utils.profiling_and_tracing.tracer as tr
 import time
 from mpi4py import MPI
 
diff --git a/hydragnn/utils/__init__.py b/hydragnn/utils/__init__.py
deleted file mode 100644
index c1e23d2e1..000000000
--- a/hydragnn/utils/__init__.py
+++ /dev/null
@@ -1,40 +0,0 @@
-from .print_utils import print_distributed, iterate_tqdm, setup_log
-from .distributed import (
-    get_comm_size_and_rank,
-    get_device_list,
-    get_device,
-    get_device_name,
-    get_device_from_name,
-    is_model_distributed,
-    get_distributed_model,
-    setup_ddp,
-    nsplit,
-    comm_reduce,
-)
-from .model import (
-    save_model,
-    get_summary_writer,
-    unsorted_segment_mean,
-    load_existing_model,
-    load_existing_model_config,
-    loss_function_selection,
-    tensor_divide,
-    EarlyStopping,
-)
-from .time_utils import Timer, print_timers
-from .config_utils import (
-    update_config,
-    update_config_minmax,
-    get_log_name_config,
-    save_config,
-)
-from .deephyper import (
-    master_from_host,
-    read_node_list,
-    create_ds_config,
-    read_job_node_list,
-    create_launch_command,
-)
-
-from .optimizer import select_optimizer
-from .atomicdescriptors import atomicdescriptors
diff --git a/hydragnn/utils/datasets/__init__.py b/hydragnn/utils/datasets/__init__.py
new file mode 100644
index 000000000..8f7028fba
--- /dev/null
+++ b/hydragnn/utils/datasets/__init__.py
@@ -0,0 +1,19 @@
+from .abstractbasedataset import AbstractBaseDataset
+from .abstractrawdataset import AbstractRawDataset
+from .adiosdataset import AdiosDataset, AdiosWriter
+from .cfgdataset import CFGDataset
+from .compositional_data_splitting import (
+    get_keys,
+    get_elements_list,
+    get_max_graph_size,
+    create_dictionary_from_elements_list,
+    create_dataset_categories,
+    duplicate_unique_data_samples,
+    generate_partition,
+    compositional_stratified_splitting,
+)
+from .distdataset import DistDataset
+from .lsmsdataset import LSMSDataset
+from .pickledataset import SimplePickleDataset, SimplePickleWriter
+from .serializeddataset import SerializedDataset, SerializedWriter
+from .xyzdataset import XYZDataset
diff --git a/hydragnn/utils/abstractbasedataset.py b/hydragnn/utils/datasets/abstractbasedataset.py
similarity index 87%
rename from hydragnn/utils/abstractbasedataset.py
rename to hydragnn/utils/datasets/abstractbasedataset.py
index 556ead164..7e73e0859 100644
--- a/hydragnn/utils/abstractbasedataset.py
+++ b/hydragnn/utils/datasets/abstractbasedataset.py
@@ -5,7 +5,7 @@
 
 class AbstractBaseDataset(torch.utils.data.Dataset, ABC):
     """
-    HydraGNN's base dataset. This is abstract class.
+    HydraGNN's base datasets. This is abstract class.
     """
 
     def __init__(self):
@@ -15,14 +15,14 @@ def __init__(self):
     @abstractmethod
     def get(self, idx):
         """
-        Return a dataset at idx
+        Return a datasets at idx
         """
         pass
 
     @abstractmethod
     def len(self):
         """
-        Total number of dataset.
+        Total number of datasets.
         If data is distributed, it should be the global total size.
         """
         pass
diff --git a/hydragnn/utils/abstractrawdataset.py b/hydragnn/utils/datasets/abstractrawdataset.py
similarity index 95%
rename from hydragnn/utils/abstractrawdataset.py
rename to hydragnn/utils/datasets/abstractrawdataset.py
index 657863ea4..d81d4bfd5 100644
--- a/hydragnn/utils/abstractrawdataset.py
+++ b/hydragnn/utils/datasets/abstractrawdataset.py
@@ -3,8 +3,6 @@
 import random
 
 import torch
-from torch import tensor
-from torch_geometric.data import Data
 from torch_geometric.transforms import (
     Distance,
     NormalizeRotation,
@@ -12,38 +10,31 @@
     PointPairFeatures,
 )
 
-from hydragnn.utils import nsplit, tensor_divide, comm_reduce
-from hydragnn.utils.print_utils import print_distributed, iterate_tqdm, log
+from hydragnn.utils.distributed import nsplit, comm_reduce
+from hydragnn.utils.model.model import tensor_divide
+from hydragnn.utils.print.print_utils import iterate_tqdm, log
 from hydragnn.utils.distributed import get_device
-from hydragnn.utils.abstractbasedataset import AbstractBaseDataset
-from hydragnn.preprocess.utils import (
+from hydragnn.utils.datasets.abstractbasedataset import AbstractBaseDataset
+from hydragnn.preprocess.graph_samples_checks_and_updates import (
     get_radius_graph,
     get_radius_graph_pbc,
-    get_radius_graph_config,
-    get_radius_graph_pbc_config,
 )
 from hydragnn.preprocess import (
-    update_predicted_values,
-    update_atom_features,
     stratified_sampling,
 )
 
-from sklearn.model_selection import StratifiedShuffleSplit
-
-from hydragnn.preprocess.dataset_descriptors import AtomFeatures
-
 from abc import ABC, abstractmethod
 
 
 class AbstractRawDataset(AbstractBaseDataset, ABC):
-    """Raw dataset class"""
+    """Raw datasets class"""
 
     def __init__(self, config, dist=False, sampling=None):
         super().__init__()
 
         """
         config:
-          shows the dataset path the target variables information, e.g, location and dimension, in data file
+          shows the datasets path the target variables information, e.g, location and dimension, in data file
         ###########
         dataset_list:
           list of datasets read from self.path_dictionary
@@ -215,7 +206,7 @@ def __load_raw_data(self):
 
     def __normalize_dataset(self):
 
-        """Performs the normalization on Data objects and returns the normalized dataset."""
+        """Performs the normalization on Data objects and returns the normalized datasets."""
         num_node_features = len(self.node_feature_dim)
         num_graph_features = len(self.graph_feature_dim)
 
diff --git a/hydragnn/utils/adiosdataset.py b/hydragnn/utils/datasets/adiosdataset.py
similarity index 98%
rename from hydragnn/utils/adiosdataset.py
rename to hydragnn/utils/datasets/adiosdataset.py
index c366cea86..d32661ac1 100644
--- a/hydragnn/utils/adiosdataset.py
+++ b/hydragnn/utils/datasets/adiosdataset.py
@@ -2,9 +2,8 @@
 import pickle
 import time
 import os
-import glob
 
-from .print_utils import print_distributed, log, log0, iterate_tqdm
+from hydragnn.utils.print.print_utils import log, log0, iterate_tqdm
 
 import numpy as np
 
@@ -23,10 +22,10 @@
 except ImportError:
     pass
 
-import hydragnn.utils.tracer as tr
+import hydragnn.utils.profiling_and_tracing.tracer as tr
 
-from hydragnn.utils.abstractbasedataset import AbstractBaseDataset
-from hydragnn.utils import nsplit
+from hydragnn.utils.datasets.abstractbasedataset import AbstractBaseDataset
+from hydragnn.utils.distributed import nsplit
 from hydragnn.preprocess import update_predicted_values, update_atom_features
 
 
@@ -279,7 +278,7 @@ def save(self):
 
 
 class AdiosDataset(AbstractBaseDataset):
-    """Adios dataset class"""
+    """Adios datasets class"""
 
     def __init__(
         self,
@@ -306,7 +305,7 @@ def __init__(
         comm: MPI_comm
             MPI communicator
         preload: bool, optional
-            Option to preload all the dataset into a memory
+            Option to preload all the datasets into a memory
         shmem: bool, optional
             Option to use shmem to share data between processes in the same node
         enable_cache: bool, optional
diff --git a/hydragnn/utils/cfgdataset.py b/hydragnn/utils/datasets/cfgdataset.py
similarity index 97%
rename from hydragnn/utils/cfgdataset.py
rename to hydragnn/utils/datasets/cfgdataset.py
index 5e7c59e7d..8df40217b 100644
--- a/hydragnn/utils/cfgdataset.py
+++ b/hydragnn/utils/datasets/cfgdataset.py
@@ -3,7 +3,7 @@
 
 from torch import tensor
 from torch_geometric.data import Data
-from hydragnn.utils.abstractrawdataset import AbstractRawDataset
+from hydragnn.utils.datasets.abstractrawdataset import AbstractRawDataset
 
 from ase.io.cfg import read_cfg
 
diff --git a/hydragnn/preprocess/compositional_data_splitting.py b/hydragnn/utils/datasets/compositional_data_splitting.py
similarity index 94%
rename from hydragnn/preprocess/compositional_data_splitting.py
rename to hydragnn/utils/datasets/compositional_data_splitting.py
index 574c10dcf..4805c1245 100644
--- a/hydragnn/preprocess/compositional_data_splitting.py
+++ b/hydragnn/utils/datasets/compositional_data_splitting.py
@@ -115,8 +115,8 @@ def generate_partition(
 
 
 def compositional_stratified_splitting(dataset, perc_train):
-    """Given the dataset and the percentage of data you want to extract from it, method will
-    apply stratified sampling where X is the dataset and Y is are the category values for each datapoint.
+    """Given the datasets and the percentage of data you want to extract from it, method will
+    apply stratified sampling where X is the datasets and Y is are the category values for each datapoint.
     In the case each structure contains 2 types of atoms, the category will
     be constructed as such: number of atoms of type 1 + number of atoms of type 2 * 100.
     Parameters
@@ -124,11 +124,11 @@ def compositional_stratified_splitting(dataset, perc_train):
     dataset: [Data]
         A list of Data objects representing a structure that has atoms.
     subsample_percentage: float
-        Percentage of the dataset.
+        Percentage of the datasets.
     Returns
     ----------
     [Data]
-        Subsample of the original dataset constructed using stratified sampling.
+        Subsample of the original datasets constructed using stratified sampling.
     """
     dataset_categories = create_dataset_categories(dataset)
     dataset, dataset_categories = duplicate_unique_data_samples(
diff --git a/hydragnn/utils/distdataset.py b/hydragnn/utils/datasets/distdataset.py
similarity index 95%
rename from hydragnn/utils/distdataset.py
rename to hydragnn/utils/datasets/distdataset.py
index 5732a6cd8..80c282e67 100644
--- a/hydragnn/utils/distdataset.py
+++ b/hydragnn/utils/datasets/distdataset.py
@@ -4,23 +4,23 @@
 import torch
 import torch_geometric.data
 
-from hydragnn.utils.abstractbasedataset import AbstractBaseDataset
+from hydragnn.utils.datasets.abstractbasedataset import AbstractBaseDataset
 
 try:
     import pyddstore as dds
 except ImportError:
     pass
 
-from hydragnn.utils.print_utils import log, log0
-from hydragnn.utils import nsplit
+from hydragnn.utils.print.print_utils import log0
+from hydragnn.utils.distributed import nsplit
 from hydragnn.preprocess import update_predicted_values, update_atom_features
 
-import hydragnn.utils.tracer as tr
+import hydragnn.utils.profiling_and_tracing.tracer as tr
 from tqdm import tqdm
 
 
 class DistDataset(AbstractBaseDataset):
-    """Distributed dataset class"""
+    """Distributed datasets class"""
 
     def __init__(
         self,
diff --git a/hydragnn/utils/lsmsdataset.py b/hydragnn/utils/datasets/lsmsdataset.py
similarity index 97%
rename from hydragnn/utils/lsmsdataset.py
rename to hydragnn/utils/datasets/lsmsdataset.py
index a1314938d..99a121644 100644
--- a/hydragnn/utils/lsmsdataset.py
+++ b/hydragnn/utils/datasets/lsmsdataset.py
@@ -1,6 +1,6 @@
 from torch import tensor
 from torch_geometric.data import Data
-from hydragnn.utils.abstractrawdataset import AbstractRawDataset
+from hydragnn.utils.datasets.abstractrawdataset import AbstractRawDataset
 
 
 class LSMSDataset(AbstractRawDataset):
diff --git a/hydragnn/utils/pickledataset.py b/hydragnn/utils/datasets/pickledataset.py
similarity index 95%
rename from hydragnn/utils/pickledataset.py
rename to hydragnn/utils/datasets/pickledataset.py
index 8b99f0f9d..48da3d06b 100644
--- a/hydragnn/utils/pickledataset.py
+++ b/hydragnn/utils/datasets/pickledataset.py
@@ -1,15 +1,14 @@
 import os
 import pickle
 
-import torch
 from mpi4py import MPI
 
-from .print_utils import print_distributed, log, iterate_tqdm
+from hydragnn.utils.print.print_utils import log, iterate_tqdm
 
-from hydragnn.utils.abstractbasedataset import AbstractBaseDataset
+from hydragnn.utils.datasets.abstractbasedataset import AbstractBaseDataset
 from hydragnn.preprocess import update_predicted_values, update_atom_features
 
-import hydragnn.utils.tracer as tr
+import hydragnn.utils.profiling_and_tracing.tracer as tr
 
 
 class SimplePickleDataset(AbstractBaseDataset):
@@ -119,7 +118,7 @@ def __init__(
         """
         Parameters
         ----------
-        dataset: locally owned dataset (should be iterable)
+        dataset: locally owned datasets (should be iterable)
         basedir: basedir
         label: label
         nmax: nmax in case of subdir
diff --git a/hydragnn/utils/serializeddataset.py b/hydragnn/utils/datasets/serializeddataset.py
similarity index 86%
rename from hydragnn/utils/serializeddataset.py
rename to hydragnn/utils/datasets/serializeddataset.py
index c469e3cab..70f71076f 100644
--- a/hydragnn/utils/serializeddataset.py
+++ b/hydragnn/utils/datasets/serializeddataset.py
@@ -1,10 +1,10 @@
 import os
 import pickle
 
-from .print_utils import log
+from hydragnn.utils.print.print_utils import log
 
-from hydragnn.utils import get_comm_size_and_rank
-from hydragnn.utils.abstractbasedataset import AbstractBaseDataset
+from hydragnn.utils.distributed import get_comm_size_and_rank
+from hydragnn.utils.datasets.abstractbasedataset import AbstractBaseDataset
 
 
 class SerializedDataset(AbstractBaseDataset):
@@ -15,7 +15,7 @@ def __init__(self, basedir, datasetname, label, dist=False):
         Parameters
         ----------
         basedir: basedir
-        datasetname: dataset name
+        datasetname: datasets name
         label: label
         """
         super().__init__()
@@ -62,9 +62,9 @@ def __init__(
         """
         Parameters
         ----------
-        dataset: locally owned dataset (should be iterable)
+        dataset: locally owned datasets (should be iterable)
         basedir: basedir
-        datasetname: dataset name
+        datasetname: datasets name
         label: label
         nmax: nmax in case of subdir
         minmax_node_feature: minmax_node_feature
diff --git a/hydragnn/utils/xyzdataset.py b/hydragnn/utils/datasets/xyzdataset.py
similarity index 95%
rename from hydragnn/utils/xyzdataset.py
rename to hydragnn/utils/datasets/xyzdataset.py
index b7c89be30..e3b57c29b 100644
--- a/hydragnn/utils/xyzdataset.py
+++ b/hydragnn/utils/datasets/xyzdataset.py
@@ -3,9 +3,8 @@
 
 from torch import tensor
 from torch_geometric.data import Data
-from hydragnn.utils.abstractrawdataset import AbstractRawDataset
+from hydragnn.utils.datasets.abstractrawdataset import AbstractRawDataset
 
-from ase.io.cfg import read_cfg
 from ase.io import read
 
 
diff --git a/hydragnn/utils/descriptors_and_embeddings/__init__.py b/hydragnn/utils/descriptors_and_embeddings/__init__.py
new file mode 100644
index 000000000..13b70bc6c
--- /dev/null
+++ b/hydragnn/utils/descriptors_and_embeddings/__init__.py
@@ -0,0 +1,6 @@
+from .atomicdescriptors import atomicdescriptors
+from .smiles_utils import (
+    get_node_attribute_name,
+    generate_graphdata_from_smilestr,
+    generate_graphdata_from_rdkit_molecule,
+)
diff --git a/hydragnn/utils/atomicdescriptors.py b/hydragnn/utils/descriptors_and_embeddings/atomicdescriptors.py
similarity index 100%
rename from hydragnn/utils/atomicdescriptors.py
rename to hydragnn/utils/descriptors_and_embeddings/atomicdescriptors.py
diff --git a/hydragnn/utils/smiles_utils.py b/hydragnn/utils/descriptors_and_embeddings/smiles_utils.py
similarity index 100%
rename from hydragnn/utils/smiles_utils.py
rename to hydragnn/utils/descriptors_and_embeddings/smiles_utils.py
diff --git a/hydragnn/utils/distributed/__init__.py b/hydragnn/utils/distributed/__init__.py
new file mode 100644
index 000000000..063642124
--- /dev/null
+++ b/hydragnn/utils/distributed/__init__.py
@@ -0,0 +1,16 @@
+from .distributed import (
+    get_comm_size_and_rank,
+    get_device_list,
+    get_device,
+    get_device_name,
+    get_device_from_name,
+    is_model_distributed,
+    get_distributed_model,
+    setup_ddp,
+    nsplit,
+    comm_reduce,
+    get_deepspeed_init_args,
+    init_comm_size_and_rank,
+    check_remaining,
+    print_peak_memory,
+)
diff --git a/hydragnn/utils/distributed.py b/hydragnn/utils/distributed/distributed.py
similarity index 99%
rename from hydragnn/utils/distributed.py
rename to hydragnn/utils/distributed/distributed.py
index 50c853776..1cc86d3a4 100644
--- a/hydragnn/utils/distributed.py
+++ b/hydragnn/utils/distributed/distributed.py
@@ -15,14 +15,13 @@
 import torch
 import torch.distributed as dist
 
-from .print_utils import print_distributed
+from hydragnn.utils.print.print_utils import print_distributed
 
 import psutil
 import socket
 from datetime import timedelta
 import time
 import subprocess
-from mpi4py import MPI
 
 deepspeed_available = True
 try:
diff --git a/hydragnn/utils/hpo/__init__.py b/hydragnn/utils/hpo/__init__.py
new file mode 100644
index 000000000..d0cf926bd
--- /dev/null
+++ b/hydragnn/utils/hpo/__init__.py
@@ -0,0 +1,7 @@
+from .deephyper import (
+    master_from_host,
+    read_node_list,
+    create_ds_config,
+    read_job_node_list,
+    create_launch_command,
+)
diff --git a/hydragnn/utils/deephyper.py b/hydragnn/utils/hpo/deephyper.py
similarity index 100%
rename from hydragnn/utils/deephyper.py
rename to hydragnn/utils/hpo/deephyper.py
diff --git a/hydragnn/utils/input_config_parsing/__init__.py b/hydragnn/utils/input_config_parsing/__init__.py
new file mode 100644
index 000000000..50e3e5176
--- /dev/null
+++ b/hydragnn/utils/input_config_parsing/__init__.py
@@ -0,0 +1,6 @@
+from .config_utils import (
+    update_config,
+    update_config_minmax,
+    get_log_name_config,
+    save_config,
+)
diff --git a/hydragnn/utils/config_utils.py b/hydragnn/utils/input_config_parsing/config_utils.py
similarity index 97%
rename from hydragnn/utils/config_utils.py
rename to hydragnn/utils/input_config_parsing/config_utils.py
index 80762cc95..86525c766 100644
--- a/hydragnn/utils/config_utils.py
+++ b/hydragnn/utils/input_config_parsing/config_utils.py
@@ -10,15 +10,14 @@
 ##############################################################################
 import pickle
 import os
-from hydragnn.preprocess.utils import check_if_graph_size_variable, gather_deg
-from hydragnn.utils.model import calculate_PNA_degree
-from hydragnn.utils import get_comm_size_and_rank
-import time
+from hydragnn.preprocess.graph_samples_checks_and_updates import (
+    check_if_graph_size_variable,
+    gather_deg,
+)
+from hydragnn.utils.distributed import get_comm_size_and_rank
 from copy import deepcopy
 import json
-from torch_geometric.utils import degree
 import torch
-import torch.distributed as dist
 
 
 def update_config(config, train_loader, val_loader, test_loader):
@@ -47,7 +46,7 @@ def update_config(config, train_loader, val_loader, test_loader):
 
     if config["NeuralNetwork"]["Architecture"]["model_type"] == "PNA" or "PNAPlus":
         if hasattr(train_loader.dataset, "pna_deg"):
-            ## Use max neighbours used in the dataset.
+            ## Use max neighbours used in the datasets.
             deg = torch.tensor(train_loader.dataset.pna_deg)
         else:
             deg = gather_deg(train_loader.dataset)
diff --git a/utils/lsms/__init__.py b/hydragnn/utils/lsms/__init__.py
similarity index 100%
rename from utils/lsms/__init__.py
rename to hydragnn/utils/lsms/__init__.py
diff --git a/utils/lsms/compositional_histogram_cutoff.py b/hydragnn/utils/lsms/compositional_histogram_cutoff.py
similarity index 100%
rename from utils/lsms/compositional_histogram_cutoff.py
rename to hydragnn/utils/lsms/compositional_histogram_cutoff.py
diff --git a/utils/lsms/convert_total_energy_to_formation_gibbs.py b/hydragnn/utils/lsms/convert_total_energy_to_formation_gibbs.py
similarity index 100%
rename from utils/lsms/convert_total_energy_to_formation_gibbs.py
rename to hydragnn/utils/lsms/convert_total_energy_to_formation_gibbs.py
diff --git a/hydragnn/utils/model/__init__.py b/hydragnn/utils/model/__init__.py
new file mode 100644
index 000000000..0ec38efbd
--- /dev/null
+++ b/hydragnn/utils/model/__init__.py
@@ -0,0 +1,12 @@
+from .model import (
+    activation_function_selection,
+    save_model,
+    get_summary_writer,
+    unsorted_segment_mean,
+    load_existing_model,
+    load_existing_model_config,
+    loss_function_selection,
+    tensor_divide,
+    EarlyStopping,
+    print_model,
+)
diff --git a/hydragnn/utils/model.py b/hydragnn/utils/model/model.py
similarity index 98%
rename from hydragnn/utils/model.py
rename to hydragnn/utils/model/model.py
index 67253fcae..6b6d3eb56 100644
--- a/hydragnn/utils/model.py
+++ b/hydragnn/utils/model/model.py
@@ -15,9 +15,8 @@
 import torch
 import torch.distributed as dist
 from torch.utils.tensorboard import SummaryWriter
-from torch_geometric.data import Data
 from torch_geometric.utils import degree
-from .print_utils import print_master, iterate_tqdm
+from hydragnn.utils.print.print_utils import print_master, iterate_tqdm
 
 from hydragnn.utils.distributed import (
     get_comm_size_and_rank,
@@ -123,7 +122,7 @@ def load_existing_model(
         model.load_checkpoint(os.path.join(path, model_name), model_name)
 
 
-## This function may cause OOM if dataset is too large
+## This function may cause OOM if datasets is too large
 ## to fit in a single GPU (i.e., with DDP). Use with caution.
 ## Recommend to use calculate_PNA_degree_dist
 def calculate_PNA_degree(loader, max_neighbours):
diff --git a/hydragnn/utils/optimizer/__init__.py b/hydragnn/utils/optimizer/__init__.py
new file mode 100644
index 000000000..9d9dce433
--- /dev/null
+++ b/hydragnn/utils/optimizer/__init__.py
@@ -0,0 +1 @@
+from .optimizer import select_optimizer
diff --git a/hydragnn/utils/optimizer.py b/hydragnn/utils/optimizer/optimizer.py
similarity index 97%
rename from hydragnn/utils/optimizer.py
rename to hydragnn/utils/optimizer/optimizer.py
index 6950e1146..c4f6e8bef 100644
--- a/hydragnn/utils/optimizer.py
+++ b/hydragnn/utils/optimizer/optimizer.py
@@ -1,5 +1,5 @@
 import torch
-from .distributed import get_device_name
+from hydragnn.utils.distributed import get_device_name
 from torch.distributed.optim import ZeroRedundancyOptimizer
 
 deepspeed_available = True
@@ -70,7 +70,7 @@ def select_zero_redundancy_optimizer(model, config):
     elif config["type"] == "Adamax":
         optimizer = ZeroRedundancyOptimizer(
             model.parameters(),
-            optimizer_class=torch.optim.Adagrad,
+            optimizer_class=torch.optim.Adamax,
             lr=config["learning_rate"],
         )
     elif config["type"] == "AdamW":
diff --git a/hydragnn/utils/print/__init__.py b/hydragnn/utils/print/__init__.py
new file mode 100644
index 000000000..8093611dc
--- /dev/null
+++ b/hydragnn/utils/print/__init__.py
@@ -0,0 +1 @@
+from .print_utils import print_distributed, iterate_tqdm, setup_log
diff --git a/hydragnn/utils/print_utils.py b/hydragnn/utils/print/print_utils.py
similarity index 95%
rename from hydragnn/utils/print_utils.py
rename to hydragnn/utils/print/print_utils.py
index f01facaf8..eb7329ecd 100644
--- a/hydragnn/utils/print_utils.py
+++ b/hydragnn/utils/print/print_utils.py
@@ -64,7 +64,7 @@ def setup_log(prefix):
     """
     Setup logging to print messages for both screen and file.
     """
-    from .distributed import init_comm_size_and_rank
+    from hydragnn.utils.distributed import init_comm_size_and_rank
 
     world_size, world_rank = init_comm_size_and_rank()
 
@@ -100,7 +100,7 @@ def log(*args, sep=" ", rank=None):
     if rank is None:
         logger.info(sep.join(map(str, args)))
     else:
-        from .distributed import init_comm_size_and_rank
+        from hydragnn.utils.distributed import init_comm_size_and_rank
 
         world_size, world_rank = init_comm_size_and_rank()
         if rank == world_rank:
diff --git a/hydragnn/utils/profiling_and_tracing/__init__.py b/hydragnn/utils/profiling_and_tracing/__init__.py
new file mode 100644
index 000000000..25c9d8399
--- /dev/null
+++ b/hydragnn/utils/profiling_and_tracing/__init__.py
@@ -0,0 +1 @@
+from .time_utils import print_timers
diff --git a/hydragnn/utils/gptl4py_dummy.py b/hydragnn/utils/profiling_and_tracing/gptl4py_dummy.py
similarity index 97%
rename from hydragnn/utils/gptl4py_dummy.py
rename to hydragnn/utils/profiling_and_tracing/gptl4py_dummy.py
index 67cafe45f..4f8fd91e5 100644
--- a/hydragnn/utils/gptl4py_dummy.py
+++ b/hydragnn/utils/profiling_and_tracing/gptl4py_dummy.py
@@ -12,7 +12,6 @@
 from __future__ import absolute_import
 from functools import wraps
 from contextlib import contextmanager
-import torch.cuda.nvtx as nvtx
 
 
 def initialize():
diff --git a/hydragnn/utils/profile.py b/hydragnn/utils/profiling_and_tracing/profile.py
similarity index 96%
rename from hydragnn/utils/profile.py
rename to hydragnn/utils/profiling_and_tracing/profile.py
index a0e113324..0e15027c8 100644
--- a/hydragnn/utils/profile.py
+++ b/hydragnn/utils/profiling_and_tracing/profile.py
@@ -1,7 +1,7 @@
 import torch
 import contextlib
 from unittest.mock import MagicMock
-from torch.profiler import profile, record_function, ProfilerActivity
+from torch.profiler import ProfilerActivity
 
 from hydragnn.utils.distributed import get_device_name
 
diff --git a/hydragnn/utils/time_utils.py b/hydragnn/utils/profiling_and_tracing/time_utils.py
similarity index 97%
rename from hydragnn/utils/time_utils.py
rename to hydragnn/utils/profiling_and_tracing/time_utils.py
index f30bb9b11..ddd7b0251 100644
--- a/hydragnn/utils/time_utils.py
+++ b/hydragnn/utils/profiling_and_tracing/time_utils.py
@@ -11,8 +11,8 @@
 
 import time
 import torch
-from .distributed import get_comm_size_and_rank, get_device
-from .print_utils import print_distributed
+from hydragnn.utils.distributed import get_comm_size_and_rank, get_device
+from hydragnn.utils.print.print_utils import print_distributed
 
 
 class TimerError(Exception):
diff --git a/hydragnn/utils/tracer.py b/hydragnn/utils/profiling_and_tracing/tracer.py
similarity index 98%
rename from hydragnn/utils/tracer.py
rename to hydragnn/utils/profiling_and_tracing/tracer.py
index df2b5285f..a156f9cf6 100644
--- a/hydragnn/utils/tracer.py
+++ b/hydragnn/utils/profiling_and_tracing/tracer.py
@@ -6,10 +6,6 @@
 from functools import wraps
 from contextlib import contextmanager
 
-import os
-import sys
-from collections import OrderedDict
-
 from abc import ABC, abstractmethod
 import torch
 from mpi4py import MPI
diff --git a/tests/test_atomicdescriptors.py b/tests/test_atomicdescriptors.py
index 0dc96ed94..3d98b2304 100644
--- a/tests/test_atomicdescriptors.py
+++ b/tests/test_atomicdescriptors.py
@@ -18,7 +18,9 @@
 @pytest.mark.mpi_skip()
 def pytest_atomicdescriptors():
     file_path = os.path.join(
-        os.path.dirname(__file__), "..", "hydragnn/utils/atomicdescriptors.py"
+        os.path.dirname(__file__),
+        "..",
+        "hydragnn/utils/descriptors_and_embeddings/atomicdescriptors.py",
     )
     return_code = subprocess.call(["python", file_path])
 
diff --git a/tests/test_datasetclass_inheritance.py b/tests/test_datasetclass_inheritance.py
index ba2e76e28..a742e0eca 100644
--- a/tests/test_datasetclass_inheritance.py
+++ b/tests/test_datasetclass_inheritance.py
@@ -20,10 +20,13 @@
 
 import hydragnn, tests
 
-from hydragnn.utils.config_utils import get_log_name_config
+from hydragnn.utils.input_config_parsing.config_utils import get_log_name_config
 from hydragnn.utils.model import print_model
-from hydragnn.utils.lsmsdataset import LSMSDataset
-from hydragnn.utils.serializeddataset import SerializedWriter, SerializedDataset
+from hydragnn.utils.datasets.lsmsdataset import LSMSDataset
+from hydragnn.utils.datasets.serializeddataset import (
+    SerializedWriter,
+    SerializedDataset,
+)
 
 from hydragnn.preprocess.load_data import split_dataset
 
diff --git a/tests/test_deepspeed.py b/tests/test_deepspeed.py
index 2cd89b4e9..f111b282c 100644
--- a/tests/test_deepspeed.py
+++ b/tests/test_deepspeed.py
@@ -1,9 +1,5 @@
-import os, json
 import pytest
 
-import torch
-import random
-import hydragnn
 from tests.test_graphs import unittest_train_model
 
 
diff --git a/tests/test_enthalpy.py b/tests/test_enthalpy.py
index 4fd7ac04c..45a98293c 100644
--- a/tests/test_enthalpy.py
+++ b/tests/test_enthalpy.py
@@ -11,9 +11,9 @@
 
 import os
 import numpy as np
-import hydragnn, tests
+import tests
 import pytest
-from utils.lsms import (
+from hydragnn.utils.lsms import (
     convert_raw_data_energy_to_gibbs,
 )
 
diff --git a/tests/test_graphs.py b/tests/test_graphs.py
index 6409a77d4..6222d707b 100755
--- a/tests/test_graphs.py
+++ b/tests/test_graphs.py
@@ -18,7 +18,7 @@
 import shutil
 
 import hydragnn, tests
-from hydragnn.utils.config_utils import merge_config
+from hydragnn.utils.input_config_parsing.config_utils import merge_config
 
 
 # Main unit test function called by pytest wrappers.
@@ -30,7 +30,7 @@ def unittest_train_model(
     use_deepspeed=False,
     overwrite_config=None,
 ):
-    world_size, rank = hydragnn.utils.get_comm_size_and_rank()
+    world_size, rank = hydragnn.utils.distributed.get_comm_size_and_rank()
 
     os.environ["SERIALIZED_DATA_PATH"] = os.getcwd()
 
@@ -167,7 +167,7 @@ def unittest_train_model(
             + " < "
             + str(thresholds[model_type][0])
         )
-        hydragnn.utils.print_distributed(verbosity, "head: " + error_str)
+        hydragnn.utils.print.print_distributed(verbosity, "head: " + error_str)
         assert (
             error_head_mse < thresholds[model_type][0]
         ), "Head RMSE checking failed for " + str(ihead)
@@ -188,7 +188,7 @@ def unittest_train_model(
 
     # Check RMSE error
     error_str = str("{:.6f}".format(error)) + " < " + str(thresholds[model_type][0])
-    hydragnn.utils.print_distributed(verbosity, "total: " + error_str)
+    hydragnn.utils.print.print_distributed(verbosity, "total: " + error_str)
     assert error < thresholds[model_type][0], "Total RMSE checking failed!" + str(error)
 
 
diff --git a/tests/test_loss_and_activation_functions.py b/tests/test_loss_and_activation_functions.py
index 22ce821a3..e0a73e739 100644
--- a/tests/test_loss_and_activation_functions.py
+++ b/tests/test_loss_and_activation_functions.py
@@ -22,7 +22,7 @@
 def unittest_loss_and_activation_functions(
     activation_function_type, loss_function_type, ci_input, overwrite_data=False
 ):
-    world_size, rank = hydragnn.utils.get_comm_size_and_rank()
+    world_size, rank = hydragnn.utils.distributed.get_comm_size_and_rank()
 
     os.environ["SERIALIZED_DATA_PATH"] = os.getcwd()
 
diff --git a/tests/test_model_loadpred.py b/tests/test_model_loadpred.py
index b301962a2..a8d650b43 100755
--- a/tests/test_model_loadpred.py
+++ b/tests/test_model_loadpred.py
@@ -29,9 +29,11 @@ def unittest_model_prediction(config):
         verbosity=config["Verbosity"]["level"],
     )
 
-    model = hydragnn.utils.get_distributed_model(model, config["Verbosity"]["level"])
+    model = hydragnn.utils.distributed.get_distributed_model(
+        model, config["Verbosity"]["level"]
+    )
 
-    log_name = hydragnn.utils.config_utils.get_log_name_config(config)
+    log_name = hydragnn.utils.input_config_parsing.get_log_name_config(config)
     hydragnn.utils.model.load_existing_model(model, log_name)
 
     model.eval()
@@ -66,9 +68,11 @@ def pytest_model_loadpred():
         config = json.load(f)
     config["NeuralNetwork"]["Architecture"]["model_type"] = model_type
     # get the directory of trained model
-    log_name = hydragnn.utils.config_utils.get_log_name_config(config)
+    log_name = hydragnn.utils.input_config_parsing.config_utils.get_log_name_config(
+        config
+    )
     modelfile = os.path.join("./logs/", log_name, log_name + ".pk")
-    # check if pretrained model and pkl dataset files exists
+    # check if pretrained model and pkl datasets files exists
     case_exist = True
     config_file = os.path.join("./logs/", log_name, "config.json")
     if not (os.path.isfile(modelfile) and os.path.isfile(config_file)):
@@ -79,7 +83,7 @@ def pytest_model_loadpred():
             config = json.load(f)
         for dataset_name, raw_data_path in config["Dataset"]["path"].items():
             if not os.path.isfile(raw_data_path):
-                print(dataset_name, "dataset not found: ", raw_data_path)
+                print(dataset_name, "datasets not found: ", raw_data_path)
                 case_exist = False
                 break
     if not case_exist:
diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
index 656067c46..c70d9de29 100644
--- a/tests/test_optimizer.py
+++ b/tests/test_optimizer.py
@@ -15,13 +15,12 @@
 import shutil
 
 import hydragnn, tests
-from hydragnn.utils.distributed import get_device_name
 
 
 # Loss function unit test called by pytest wrappers.
 # Note the intent of this test is to make sure all interfaces work - it does not assert anything
 def unittest_optimizers(optimizer_type, use_zero, ci_input, overwrite_data=False):
-    world_size, rank = hydragnn.utils.get_comm_size_and_rank()
+    world_size, rank = hydragnn.utils.distributed.get_comm_size_and_rank()
 
     os.environ["SERIALIZED_DATA_PATH"] = os.getcwd()
 
diff --git a/tests/test_periodic_boundary_conditions.py b/tests/test_periodic_boundary_conditions.py
index 43d92b46c..a81e0b9f6 100644
--- a/tests/test_periodic_boundary_conditions.py
+++ b/tests/test_periodic_boundary_conditions.py
@@ -9,12 +9,11 @@
 # SPDX-License-Identifier: BSD-3-Clause                                      #
 ##############################################################################
 
-import sys, os, json, numpy as np
-import pytest
+import json, numpy as np
 
 import torch
 from torch_geometric.data import Data
-from hydragnn.preprocess.utils import (
+from hydragnn.preprocess.graph_samples_checks_and_updates import (
     get_radius_graph_config,
     get_radius_graph_pbc_config,
 )
diff --git a/tests/test_rotational_invariance.py b/tests/test_rotational_invariance.py
index 6dfcf377b..f7ac970a2 100644
--- a/tests/test_rotational_invariance.py
+++ b/tests/test_rotational_invariance.py
@@ -9,15 +9,15 @@
 # SPDX-License-Identifier: BSD-3-Clause                                      #
 ##############################################################################
 
-import sys, os, json
+import json
 import pytest
 
 import torch
 from torch_geometric.data import Data
 from torch_geometric.transforms import Distance, NormalizeRotation
-from hydragnn.preprocess.utils import get_radius_graph_config
+from hydragnn.preprocess.graph_samples_checks_and_updates import get_radius_graph_config
 
-from hydragnn.preprocess.utils import (
+from hydragnn.preprocess.graph_samples_checks_and_updates import (
     check_data_samples_equivalence,
 )
 
diff --git a/utils/__init__.py b/utils/__init__.py
deleted file mode 100644
index 69448aef7..000000000
--- a/utils/__init__.py
+++ /dev/null
@@ -1 +0,0 @@
-from . import lsms