Commenting MACE utils in torch geometric (draft 2)

ORNL · Sep 27, 2024 · cde4d12 · cde4d12
1 parent 2cfd8d6
commit cde4d12
Show file tree

Hide file tree

Showing 9 changed files with 1,113 additions and 1,122 deletions.
diff --git a/hydragnn/utils/mace_utils/modules/__init__.py b/hydragnn/utils/mace_utils/modules/__init__.py
@@ -3,8 +3,6 @@
 import torch
 
 from .blocks import (
-    # AgnosticNonlinearInteractionBlock,
-    # AgnosticResidualNonlinearInteractionBlock,
     AtomicEnergiesBlock,
     EquivariantProductBasisBlock,
     InteractionBlock,
@@ -15,22 +13,14 @@
     NonLinearReadoutBlock,
     RadialEmbeddingBlock,
     RealAgnosticAttResidualInteractionBlock,
-    # RealAgnosticInteractionBlock,
-    # RealAgnosticResidualInteractionBlock,
-    # ResidualElementDependentInteractionBlock,
     ScaleShiftBlock,
 )
 
 from .radial import BesselBasis, GaussianBasis, PolynomialCutoff, ZBLBasis
 from .symmetric_contraction import SymmetricContraction
 
 interaction_classes: Dict[str, Type[InteractionBlock]] = {
-    # "AgnosticNonlinearInteractionBlock": AgnosticNonlinearInteractionBlock,
-    # "ResidualElementDependentInteractionBlock": ResidualElementDependentInteractionBlock,
-    # "AgnosticResidualNonlinearInteractionBlock": AgnosticResidualNonlinearInteractionBlock,
-    # "RealAgnosticResidualInteractionBlock": RealAgnosticResidualInteractionBlock,
     "RealAgnosticAttResidualInteractionBlock": RealAgnosticAttResidualInteractionBlock,
-    # "RealAgnosticInteractionBlock": RealAgnosticInteractionBlock,
 }
 
 gate_dict: Dict[str, Optional[Callable]] = {

diff --git a/hydragnn/utils/mace_utils/modules/utils.py b/hydragnn/utils/mace_utils/modules/utils.py
@@ -15,7 +15,7 @@
 
 from hydragnn.utils.mace_utils.tools import to_numpy
 from hydragnn.utils.mace_utils.tools.scatter import scatter_sum
-from hydragnn.utils.mace_utils.tools.torch_geometric.batch import Batch
+# from hydragnn.utils.mace_utils.tools.torch_geometric.batch import Batch
 
 from .blocks import AtomicEnergiesBlock
 

diff --git a/hydragnn/utils/mace_utils/tools/torch_geometric/__init__.py b/hydragnn/utils/mace_utils/tools/torch_geometric/__init__.py
@@ -1,7 +1,8 @@
-from .batch import Batch
-from .data import Data
-from .dataloader import DataLoader
-from .dataset import Dataset
-from .seed import seed_everything
+# from .batch import Batch
+# from .data import Data
+# from .dataloader import DataLoader
+# from .dataset import Dataset
+# from .seed import seed_everything
 
-__all__ = ["Batch", "Data", "Dataset", "DataLoader", "seed_everything"]
+# __all__ = ["Batch", "Data", "Dataset", "DataLoader", "seed_everything"]
+# __all__ = ["Data", "Dataset", "DataLoader", "seed_everything"]