model/equiformer_v3.py

"""Equiformer V3 model adapted to the OneScience FairChem runtime."""

import math

import torch
from torch.utils.checkpoint import checkpoint

from onescience.modules.func_utils.uma_graph.mixin import GraphModelMixin
from onescience.modules.layer.equiformer_v3.edge_rot_mat import init_edge_rot_mat
from onescience.modules.layer.equiformer_v3.envelope import PolynomialEnvelope
from onescience.modules.layer.equiformer_v3.input_block import EdgeDegreeEmbedding
from onescience.modules.layer.equiformer_v3.layer_norm import (
    EquivariantLayerNorm,
    EquivariantMergeLayerNorm,
    EquivariantSeparableLayerNorm,
    RMSNorm,
    get_normalization_layer,
)
from onescience.modules.layer.equiformer_v3.output_block import (
    EquivariantGraphAttentionStressHead,
    FeedForwardNetworkStressHead,
    ScalarFeedForwardNetwork,
)
from onescience.modules.layer.equiformer_v3.radial_function import (
    GaussianSmearing,
    RadialFunction,
)
from onescience.modules.layer.equiformer_v3.so3 import (
    SO3Rotation,
    SO3Linear,
)
from onescience.modules.layer.equiformer_v3.transformer_block import (
    EquivariantGraphAttention,
    FeedForwardNetwork,
    TransBlockV3,
)
from onescience.modules.layer.equiformer_v3.wigner import configure_jd_path
from onescience.utils.uma.common.registry import registry
from onescience.utils.uma.common.utils import conditional_grad


# Statistics of IS2RE 100K
_AVG_NUM_NODES = 77.81317
_AVG_DEGREE = 23.395238876342773    # IS2RE: 100k, max_radius = 5, max_neighbors = 100

_NORM_SCALE_NODES = math.sqrt(_AVG_NUM_NODES)   # 8.82117735906041
_NORM_SCALE_DEGREE = math.sqrt(_AVG_DEGREE)     # 4.836862503353054


@registry.register_model("equiformer_v3")
class EquiformerV3_OC(torch.nn.Module, GraphModelMixin):
    """
    Args:
        use_pbc (bool):         Use periodic boundary conditions
        use_pbc_single (bool, optional):    Process batch PBC graphs one at a time.
                                            This argument is set to True when training on MPTraj.
        otf_graph (bool):       Compute graph On The Fly (OTF)

        regress_forces (bool):  Compute forces
        regress_stress (bool):  Compute stress
        direct_prediction (bool):   Whether to use direct methods to predict forces and stress

        max_neighbors (int):    Maximum number of neighbors per atom
        max_radius (float):     Maximum distance between nieghboring atoms in Angstroms
        num_radial_basis (int): Number of radial basis functions
        max_num_elements (int): Maximum atomic number

        num_layers (int):           Number of layers in the GNN
        num_channels (int):         Number of channels in node embeddings
        attn_hidden_channels (int): Number of hidden channels in equivariant graph attention
        num_heads (int):            Number of attention heads
        attn_alpha_channels (int):  Number of channels for alpha vector in each attention head
        attn_value_channels (int):  Number of channels for value vector in each attention head
        ffn_hidden_channels (int):  Number of hidden channels in feedforward network
        norm_type (str):            Type of normalization layer
                                    (['sep_layer_norm', 'merge_layer_norm',
                                    'merge_layer_norm_attn_rms_norm', 'merge_rms_norm'])

        lmax (int):                 Maximum degrees (l)
        mmax (int):                 Maximum order (m)
        attn_grid_resolution_list (list:int):
                                    Grid resolution list in class `SO3Grid` in attention
        ffn_grid_resolution_list (list:int):
                                    Grid resolution list in class `SO3Grid` in feedforward network

        edge_channels (int):                Number of channels for edge-wise invariant features
        use_atom_edge_embedding (bool):     Whether to use atomic embedding along with relative distance for edge scalar features
        use_envelope (bool):        Whether to apply an envelope function to attention

        attn_activation (str):      Type of activation function in equivariant graph attention
        use_attn_renorm (bool):     Whether to re-normalize attention weights
        use_add_merge (bool):       Default: False
                                    If True, use addition to merge the source/target node features instead of concat,
                                    which can save 2x compute when rotating with Wigner-D matrices.
        use_rad_l_parametrization (bool):
                                    Default: True
                                    If True, all the m components within the same type-L vector will share the same
                                    weight from the radial function.
        softcap (float):            Default: None
                                    If not None, use soft capping to limit the range of attention logits to
                                    [- `softcap`, + `softcap`].
        attn_eps (float):           Default: 1e-16
                                    Epsilon value used in the softmax operation of attention
        ffn_activation (str):       Type of activation function for feedforward network
        use_grid_mlp (bool):        If `True`, use projecting to grids and performing MLPs for FFNs.

        use_gate_force_head (bool): If `True`, use `GateActivation` in the equivariant attention of the force prediction head.

        alpha_drop (float):         Dropout rate for the hidden features in non-linear MLP attention
        attn_mask_rate (float):     Mask rate for neighbors considered in attention
        attn_weights_drop (float):  Dropout rate for attention weights
        value_drop (float):         Dropout rate for the hidden features in non-linear value vectors
        drop_path_rate (float):     Drop path rate
        proj_drop (float):          Dropout rate for outputs of attention and FFN in Transformer blocks
        ffn_drop (float):           Dropout rate for the hidden features in FFN
        use_head_reg (bool):        Whether to apply regularization to output head (dummy argument for backend compatibility)

        gradient_checkpointing_block_list (list):
                                    A list indicating which block we apply gradient/activation checkpointing to save memory.

        avg_num_nodes (float):   Normalization factor for sum aggregation over nodes
        avg_degree (float):      Normalization factor for sum aggregation over edges

        enforce_max_neighbors_strictly (bool):      When edges are subselected based on the `max_neighbors` arg, arbitrarily select amongst equidistant / degenerate edges to have exactly the correct number.
    """
    def __init__(
        self,

        use_pbc=True,
        use_pbc_single=False,
        otf_graph=True,

        regress_forces=True,
        regress_stress=False,
        direct_prediction=True,

        max_neighbors=20,
        max_radius=12.0,
        num_radial_basis=600,
        max_num_elements=128,

        num_layers=12,
        num_channels=128,
        attn_hidden_channels=64,
        num_heads=8,
        attn_alpha_channels=32,
        attn_value_channels=16,
        ffn_hidden_channels=128,
        norm_type='merge_layer_norm',

        lmax=6,
        mmax=2,
        attn_grid_resolution_list=[20, 8],
        ffn_grid_resolution_list=[20, 20],

        edge_channels=128,
        use_atom_edge_embedding=True,
        use_envelope=True,

        attn_activation='sep-merge_gates2_swiglu',
        use_attn_renorm=True,
        use_add_merge=False,
        use_rad_l_parametrization=True,
        softcap=None,
        attn_eps=1e-16,
        ffn_activation='sep-merge_gates2_swiglu',
        use_grid_mlp=True,

        use_gate_force_head=True,

        alpha_drop=0.0,
        attn_mask_rate=0.0,
        attn_weights_drop=0.1,
        value_drop=0.0,
        drop_path_rate=0.05,
        proj_drop=0.0,
        ffn_drop=0.0,
        use_head_reg=False,

        gradient_checkpointing_block_list=None,

        avg_num_nodes=_AVG_NUM_NODES,
        avg_degree=_AVG_DEGREE,

        enforce_max_neighbors_strictly=True,
        jd_path=None,
    ):
        super().__init__()

        configure_jd_path(jd_path)

        self.use_pbc = use_pbc
        self.use_pbc_single = use_pbc_single
        self.otf_graph = otf_graph

        self.regress_forces = regress_forces
        self.regress_stress = regress_stress
        self.direct_prediction = direct_prediction
        self.direct_forces = direct_prediction

        self.max_neighbors = max_neighbors
        self.max_radius = max_radius
        self.cutoff = max_radius
        self.num_radial_basis = num_radial_basis
        self.max_num_elements = max_num_elements

        self.num_layers = num_layers
        self.num_channels = num_channels
        self.attn_hidden_channels = attn_hidden_channels
        self.num_heads = num_heads
        self.attn_alpha_channels = attn_alpha_channels
        self.attn_value_channels = attn_value_channels
        self.ffn_hidden_channels = ffn_hidden_channels
        self.norm_type = norm_type

        self.lmax = lmax
        self.mmax = mmax
        self.attn_grid_resolution_list = attn_grid_resolution_list
        self.ffn_grid_resolution_list = ffn_grid_resolution_list

        self.edge_channels = edge_channels
        self.use_atom_edge_embedding = use_atom_edge_embedding
        self.use_envelope = use_envelope

        self.attn_activation = attn_activation
        self.use_attn_renorm = use_attn_renorm
        self.use_add_merge = use_add_merge
        self.use_rad_l_parametrization = use_rad_l_parametrization
        self.softcap = softcap
        self.attn_eps = attn_eps
        self.ffn_activation = ffn_activation
        self.use_grid_mlp = use_grid_mlp

        self.use_gate_force_head = use_gate_force_head

        self.alpha_drop = alpha_drop
        self.attn_mask_rate = attn_mask_rate
        self.attn_weights_drop = attn_weights_drop
        self.value_drop = value_drop
        self.drop_path_rate = drop_path_rate
        self.proj_drop = proj_drop
        self.ffn_drop = ffn_drop
        self.use_head_reg = use_head_reg

        self.gradient_checkpointing_block_list = gradient_checkpointing_block_list
        if self.gradient_checkpointing_block_list is not None:
            assert len(self.gradient_checkpointing_block_list) == self.num_layers
        else:
            self.gradient_checkpointing_block_list = [0] * self.num_layers

        self.avg_num_nodes = avg_num_nodes
        self.avg_degree = avg_degree

        self.enforce_max_neighbors_strictly = enforce_max_neighbors_strictly

        # Atom-type embedding
        self.sphere_embedding = torch.nn.Embedding(self.max_num_elements, self.num_channels)

        # Radial basis function
        self.distance_expansion = GaussianSmearing(
            0.0,
            self.cutoff,
            self.num_radial_basis,
            2.0,
        )
        edge_input_channels = int(self.distance_expansion.num_output)

        # The sizes of radial functions (input channels and 2 hidden channels)
        self.edge_channels_list = [edge_input_channels] + [self.edge_channels] * 2

        # Envelope function
        self.envelope_func = PolynomialEnvelope(
            cutoff=self.cutoff,
            exponent=5
        ) if self.use_envelope else None

        # Computing Wigner-D matrices
        self.so3_rotation = SO3Rotation(self.lmax, self.mmax, use_rotation_mask=(not self.direct_prediction))

        # Edge-degree embedding
        self.edge_degree_embedding = EdgeDegreeEmbedding(
            num_channels=self.num_channels,
            lmax=self.lmax,
            mmax=self.mmax,
            so3_rotation=self.so3_rotation,
            max_num_elements=self.max_num_elements,
            edge_channels_list=self.edge_channels_list,
            use_atom_edge_embedding=self.use_atom_edge_embedding,
            rescale_factor=self.avg_degree
        )

        # Transformer block
        self.blocks = torch.nn.ModuleList()
        for i in range(self.num_layers):
            if self.gradient_checkpointing_block_list[i] == 1:
                attn_activation = self.attn_activation.replace('_mem', '')
                ffn_activation  = self.ffn_activation.replace('_mem', '')
            else:
                attn_activation = self.attn_activation
                ffn_activation  = self.ffn_activation
            block_config_dict = dict(
                num_in_channels=self.num_channels,
                attn_hidden_channels=self.attn_hidden_channels,
                num_heads=self.num_heads,
                attn_alpha_channels=self.attn_alpha_channels,
                attn_value_channels=self.attn_value_channels,
                ffn_hidden_channels=self.ffn_hidden_channels,
                num_out_channels=self.num_channels,
                lmax=self.lmax,
                mmax=self.mmax,
                so3_rotation=self.so3_rotation,
                attn_grid_resolution_list=self.attn_grid_resolution_list,
                ffn_grid_resolution_list=self.ffn_grid_resolution_list,
                max_num_elements=self.max_num_elements,
                edge_channels_list=self.edge_channels_list,
                use_atom_edge_embedding=self.use_atom_edge_embedding,
                attn_activation=attn_activation,
                use_attn_renorm=self.use_attn_renorm,
                use_add_merge=self.use_add_merge,
                use_rad_l_parametrization=self.use_rad_l_parametrization,
                softcap=self.softcap,
                attn_eps=self.attn_eps,
                ffn_activation=ffn_activation,
                use_grid_mlp=self.use_grid_mlp,
                norm_type=self.norm_type,
                alpha_drop=self.alpha_drop,
                attn_mask_rate=self.attn_mask_rate,
                attn_weights_drop=attn_weights_drop,
                value_drop=self.value_drop,
                drop_path_rate=self.drop_path_rate,
                proj_drop=self.proj_drop,
                ffn_drop=self.ffn_drop
            )
            block_class = TransBlockV3
            self.blocks.append(block_class(**block_config_dict))

        # Output blocks for energy and forces (and optionally stress)
        self.norm = get_normalization_layer(
            self.norm_type,
            lmax=self.lmax,
            num_channels=self.num_channels
        )
        self.energy_block = ScalarFeedForwardNetwork(
            num_in_channels=self.num_channels,
            num_hidden_channels=self.ffn_hidden_channels,
            num_out_channels=1,
            dropout=0.0
        )
        if self.direct_prediction:
            if self.regress_forces:
                self.force_block = EquivariantGraphAttention(
                    num_in_channels=self.num_channels,
                    num_hidden_channels=self.attn_hidden_channels,
                    num_heads=self.num_heads,
                    attn_alpha_channels=self.attn_alpha_channels,
                    attn_value_channels=self.attn_value_channels,
                    num_out_channels=1,
                    lmax=self.lmax,
                    mmax=self.mmax,
                    so3_rotation=self.so3_rotation,
                    grid_resolution_list=self.attn_grid_resolution_list,
                    max_num_elements=self.max_num_elements,
                    edge_channels_list=self.edge_channels_list,
                    use_atom_edge_embedding=self.use_atom_edge_embedding,
                    activation=('sep_s2' if not self.use_gate_force_head else 'gate'),
                    use_attn_renorm=self.use_attn_renorm,
                    use_add_merge=self.use_add_merge,
                    use_rad_l_parametrization=self.use_rad_l_parametrization,
                    softcap=self.softcap,
                    eps=self.attn_eps,
                    alpha_drop=0.0,
                    attn_mask_rate=0.0,
                    attn_weights_drop=0.0,
                    value_drop=0.0
                )
                if 'rms_norm' in norm_type:
                    if self.force_block.alpha_norm is not None:
                        del self.force_block.alpha_norm
                        self.force_block.alpha_norm = RMSNorm(self.attn_alpha_channels)
            if self.regress_stress:
                self.stress_block = FeedForwardNetworkStressHead(
                    num_in_channels=self.num_channels,
                    num_hidden_channels=self.ffn_hidden_channels,
                    num_out_channels=1,
                    lmax=self.lmax,
                    mmax=self.mmax,
                    grid_resolution_list=self.ffn_grid_resolution_list,
                    activation='gate',
                    use_grid_mlp=False,
                    dropout=0.0,
                )

        self.apply(self._init_weights)


    def _forward_edge(
        self,
        edge_distance,
        edge_distance_vec
    ):
        # Compute 3x3 rotation matrix per edge
        edge_rot_mat = self._init_edge_rot_mat(edge_distance_vec)

        # Compute Wigner-D matrices
        self.so3_rotation.set_wigner(edge_rot_mat)

        # Envelope function
        edge_envelope_weight = self.envelope_func(edge_distance) if self.envelope_func is not None else None

        # Radial basis function
        edge_distance = self.distance_expansion(edge_distance)

        return edge_distance, edge_envelope_weight


    def _forward_embedding(
        self,
        atomic_numbers,
        edge_distance,
        edge_index,
        edge_envelope_weight
    ):
        num_atoms = len(atomic_numbers)

        # Initialize node embedding
        x = torch.zeros(
            (
                num_atoms,
                ((self.lmax + 1) ** 2),
                self.num_channels
            ),
            device=self.device,
            dtype=self.dtype
        )

        # Atom-type embedding
        atom_embedding = self.sphere_embedding(atomic_numbers)
        x[:, 0, :] = atom_embedding

        # Edge-degree embedding
        edge_degree = self.edge_degree_embedding(
            atomic_numbers,
            edge_distance,
            edge_index,
            edge_envelope_weight
        )
        x = x + edge_degree

        return x


    def _forward_blocks(
        self,
        x,
        source_atomic_numbers,
        target_atomic_numbers,
        edge_distance,
        edge_index,
        edge_envelope_weight,
        batch
    ):
        # Transformer blocks
        for i in range(self.num_layers):
            if self.gradient_checkpointing_block_list[i] == 0:
                x = self.blocks[i](
                    x,
                    source_atomic_numbers,
                    target_atomic_numbers,
                    edge_distance,
                    edge_index,
                    edge_envelope_weight,
                    batch,     # for GraphDropPath
                )
            elif self.gradient_checkpointing_block_list[i] == 1:
                x = checkpoint(
                    self.blocks[i],
                    x,
                    source_atomic_numbers,
                    target_atomic_numbers,
                    edge_distance,
                    edge_index,
                    edge_envelope_weight,
                    batch,     # for GraphDropPath
                    use_reentrant=False
                )
            else:
                raise ValueError

        # Final layer norm
        x = self.norm(x)
        x_scalar = x.narrow(1, 0, 1)
        x_scalar = x_scalar.view(x_scalar.shape[0], self.num_channels)
        return x_scalar, x


    def _forward_direct(self, data):
        self.batch_size = len(data.natoms)
        self.dtype = data.pos.dtype
        self.device = data.pos.device

        (
            edge_index,
            edge_distance,
            edge_distance_vec,
            cell_offsets,
            _,  # cell offset distances
            neighbors,
        ) = self.generate_graph(
            data,
            enforce_max_neighbors_strictly=self.enforce_max_neighbors_strictly,
            use_pbc_single=self.use_pbc_single
        )

        atomic_numbers = data.atomic_numbers.long()
        source_atomic_numbers = atomic_numbers[edge_index[0]]
        target_atomic_numbers = atomic_numbers[edge_index[1]]

        edge_distance, edge_envelope_weight = self._forward_edge(edge_distance, edge_distance_vec)
        x = self._forward_embedding(atomic_numbers, edge_distance, edge_index, edge_envelope_weight)
        x_scalar, x = self._forward_blocks(
            x,
            source_atomic_numbers,
            target_atomic_numbers,
            edge_distance,
            edge_index,
            edge_envelope_weight,
            data.batch
        )

        outputs = {}

        # Energy prediction
        node_energy = self.energy_block(x_scalar)
        energy = torch.zeros(self.batch_size, device=node_energy.device, dtype=node_energy.dtype)
        energy.index_add_(0, data.batch, node_energy.view(-1))
        energy = energy / self.avg_num_nodes
        outputs['energy'] = energy

        # Force prediction
        if self.regress_forces:
            forces = self.force_block(
                x,
                source_atomic_numbers,
                target_atomic_numbers,
                edge_distance,
                edge_index,
                edge_envelope_weight
            )
            forces = forces.narrow(1, 1, 3)
            forces = forces.view(-1, 3)
            outputs['forces'] = forces

        # Stress Prediction
        if self.regress_stress:
            stress = self.stress_block(
                x,
                batch_size=self.batch_size,
                batch=data.batch
            )
            outputs['stress'] = stress

        return outputs


    @conditional_grad(torch.enable_grad())
    def _forward_gradient(self, data):
        """
            1.  We have additional `@conditional_grad` as the decorator since the decorator might not be compatible with
                `torch.compile()` in direct methods.
        """
        self.batch_size = len(data.natoms)
        self.dtype = data.pos.dtype
        self.device = data.pos.device

        """
            For gradient methods
        """
        displacement = None
        orig_cell = None
        if self.regress_stress and self.regress_forces:
            displacement = torch.zeros(
                (3, 3),
                dtype=self.dtype,
                device=self.device,
            )
            displacement = displacement.view(-1, 3, 3).expand(self.batch_size, 3, 3)
            displacement.requires_grad = True
            symmetric_displacement = 0.5 * (
                displacement + displacement.transpose(-1, -2)
            )

            data.pos.requires_grad = True
            data.pos = data.pos + torch.bmm(
                data.pos.unsqueeze(-2),
                torch.index_select(symmetric_displacement, 0, data.batch)
            ).squeeze(-2)

            orig_cell = data.cell
            data.cell = data.cell + torch.bmm(
                data.cell, symmetric_displacement
            )
        elif not self.regress_stress and self.regress_forces:
            data.pos.requires_grad = True

        (
            edge_index,
            edge_distance,
            edge_distance_vec,
            cell_offsets,
            _,  # cell offset distances
            neighbors,
        ) = self.generate_graph(
            data,
            enforce_max_neighbors_strictly=self.enforce_max_neighbors_strictly,
            use_pbc_single=self.use_pbc_single
        )

        atomic_numbers = data.atomic_numbers.long()
        source_atomic_numbers = atomic_numbers[edge_index[0]]
        target_atomic_numbers = atomic_numbers[edge_index[1]]

        edge_distance, edge_envelope_weight = self._forward_edge(edge_distance, edge_distance_vec)
        x = self._forward_embedding(atomic_numbers, edge_distance, edge_index, edge_envelope_weight)
        x_scalar, x = self._forward_blocks(
            x,
            source_atomic_numbers,
            target_atomic_numbers,
            edge_distance,
            edge_index,
            edge_envelope_weight,
            data.batch
        )

        outputs = {}

        # Energy prediction
        node_energy = self.energy_block(x_scalar)
        energy = torch.zeros(self.batch_size, device=node_energy.device, dtype=node_energy.dtype)
        energy.index_add_(0, data.batch, node_energy.view(-1))
        energy = energy / self.avg_num_nodes
        outputs['energy'] = energy

        if self.regress_stress and self.regress_forces:
            # Stress and forces prediction
            grads = torch.autograd.grad(
                [energy.sum()],
                [data.pos, displacement],
                create_graph=self.training,
            )
            forces = torch.neg(grads[0])
            virial = grads[1].view(-1, 3, 3)
            volume = torch.det(data.cell).abs().unsqueeze(-1)
            stress = virial / volume.view(-1, 1, 1)
            virial = torch.neg(virial)
            outputs['forces'] = forces
            outputs['stress'] = stress.view(-1, 9)
            data.cell = orig_cell

        elif not self.regress_stress and self.regress_forces:
            forces = (
                -1
                * torch.autograd.grad(
                    energy.sum(), data.pos, create_graph=self.training
                )[0]
            )
            outputs['forces'] = forces

        return outputs


    def forward(self, data):
        if self.direct_prediction:
            outputs = self._forward_direct(data)
        else:
            outputs = self._forward_gradient(data)
        return outputs


    # Initialize the edge rotation matrics
    def _init_edge_rot_mat(self, edge_distance_vec):
        return init_edge_rot_mat(edge_distance_vec, use_rotation_mask=(not self.direct_prediction))


    @property
    def num_params(self):
        return sum(p.numel() for p in self.parameters())


    def _init_weights(self, m):
        if (isinstance(m, torch.nn.Linear)
            or isinstance(m, SO3Linear)
        ):
            if m.bias is not None:
                torch.nn.init.constant_(m.bias, 0)
        elif isinstance(m, torch.nn.LayerNorm):
            torch.nn.init.constant_(m.bias, 0)
            torch.nn.init.constant_(m.weight, 1.0)
        elif (isinstance(m, RadialFunction)):
            m.apply(self._uniform_init_linear_weights)


    def _uniform_init_linear_weights(self, m):
        if isinstance(m, torch.nn.Linear):
            if m.bias is not None:
                torch.nn.init.constant_(m.bias, 0)
            std = 1 / math.sqrt(m.in_features)
            torch.nn.init.uniform_(m.weight, -std, std)


    @torch.jit.ignore
    def no_weight_decay(self):
        no_wd_list = []
        named_parameters_list = [name for name, _ in self.named_parameters()]
        for module_name, module in self.named_modules():
            if (isinstance(module, torch.nn.Embedding)
                or isinstance(module, torch.nn.Linear)
                or isinstance(module, SO3Linear)
                or isinstance(module, torch.nn.LayerNorm)
                or isinstance(module, RMSNorm)
                or isinstance(module, EquivariantLayerNorm)
                or isinstance(module, EquivariantSeparableLayerNorm)
                or isinstance(module, EquivariantMergeLayerNorm)
            ):
                for parameter_name, _ in module.named_parameters():
                    if (isinstance(module, torch.nn.Linear)
                        or isinstance(module, SO3Linear)
                    ):
                        if 'weight' in parameter_name:
                            continue
                    global_parameter_name = module_name + '.' + parameter_name
                    assert global_parameter_name in named_parameters_list
                    no_wd_list.append(global_parameter_name)
        return set(no_wd_list)


    @torch._dynamo.disable
    def generate_graph(
        self,
        data,
        cutoff=None,
        max_neighbors=None,
        use_pbc=None,
        otf_graph=None,
        enforce_max_neighbors_strictly=None,
        use_pbc_single=False,
    ):
        graph_data = super().generate_graph(
            data,
            cutoff=cutoff,
            max_neighbors=max_neighbors,
            use_pbc=use_pbc,
            otf_graph=otf_graph,
            enforce_max_neighbors_strictly=enforce_max_neighbors_strictly,
            use_pbc_single=use_pbc_single,
        )

        edge_index   = graph_data.edge_index
        edge_dist    = graph_data.edge_distance
        distance_vec = graph_data.edge_distance_vec
        cell_offsets = graph_data.cell_offsets
        cell_offset_distances = graph_data.offset_distances
        neighbors    = graph_data.neighbors

        return (
            edge_index,
            edge_dist,
            distance_vec,
            cell_offsets,
            cell_offset_distances,
            neighbors,
        )