Allegro TorchScript runtime error: size of tensor a (...) must match size of tensor b (...) at non-singleton dimension 0

[olivia614]

Hello! I’m running an Allegro/NequIP model with LAMMPS and getting a TorchScript interpreter runtime error that looks like an off-by-one mismatch between active_edges and new_latents inside Allegro:

Exception: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript, serialized code (most recent call last):
  File "code/__torch__/nequip/nn/_graph_model.py", line 29, in <forward op>
        pass
    model = self.model
    return (model).forward(new_data, )
            ~~~~~~~~~~~~~~ <--- HERE
  File "code/__torch__/nequip/nn/_rescale.py", line 21, in batch_norm
    data: Dict[str, Tensor]) -> Dict[str, Tensor]:
    model = self.model
    data0 = (model).forward(data, )
             ~~~~~~~~~~~~~~ <--- HERE
    training = self.training
    if training:
  File "code/__torch__/nequip/nn/_grad_output.py", line 71, in layer_norm
      pass
    func = self.func
    data0 = (func).forward(data, )
             ~~~~~~~~~~~~~ <--- HERE
    _17 = [torch.sum(data0["total_energy"])]
    _18 = [pos, data0["_displacement"]]
  File "code/__torch__/nequip/nn/_graph_mixin.py", line 28, in AD_logsumexp_backward
    input1 = (radial_basis).forward(input0, )
    input2 = (spharm).forward(input1, )
    input3 = (allegro).forward(input2, )
              ~~~~~~~~~~~~~~~~ <--- HERE
    input4 = (edge_eng).forward(input3, )
    input5 = (edge_eng_sum).forward(input4, )
  File "code/__torch__/allegro/nn/_allegro.py", line 178, in mean_0
    _39 = annotate(List[Optional[Tensor]], [active_edges1])
    _40 = torch.unsqueeze(torch.index(cutoff_coeffs0, _39), -1)
    new_latents2 = torch.mul(_40, new_latents1)
                   ~~~~~~~~~ <--- HERE
    latent_resnet = self.latent_resnet
    if latent_resnet:

Traceback of TorchScript, original code (most recent call last):
  File "/opt/apps/miniforge3/lib/python3.12/site-packages/nequip/nn/_graph_model.py", line 112, in <forward op>
                new_data[k] = v
        # run the model
        data = self.model(new_data)
               ~~~~~~~~~~ <--- HERE
        return data
  File "/opt/apps/miniforge3/lib/python3.12/site-packages/nequip/nn/_rescale.py", line 144, in batch_norm
    def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
        data = self.model(data)
               ~~~~~~~~~~ <--- HERE
        if self.training:
            # no scaling, but still need to promote for consistent dtype behavior
  File "/opt/apps/miniforge3/lib/python3.12/site-packages/nequip/nn/_grad_output.py", line 305, in layer_norm
    
        # Call model and get gradients
        data = self.func(data)
               ~~~~~~~~~ <--- HERE
    
        grads = torch.autograd.grad(
  File "/opt/apps/miniforge3/lib/python3.12/site-packages/nequip/nn/_graph_mixin.py", line 366, in AD_logsumexp_backward
    def forward(self, input: AtomicDataDict.Type) -> AtomicDataDict.Type:
        for module in self:
            input = module(input)
                    ~~~~~~ <--- HERE
        return input
  File "/opt/apps/miniforge3/lib/python3.12/site-packages/allegro/nn/_allegro.py", line 612, in mean_0
            torch.cat(latent_inputs_to_cat, dim=-1)[prev_mask]
        )
        new_latents = cutoff_coeffs[active_edges].unsqueeze(-1) * new_latents
                      ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
        if self.latent_resnet:
            this_layer_update_coeff = layer_update_coefficients[layer_index - 1]
RuntimeError: The size of tensor a (36675) must match the size of tensor b (36674) at non-singleton dimension 0


Abort(1) on node 0 (rank 0 in comm 0): application called MPI_Abort(MPI_COMM_WORLD, 1) - process 0

I am using:

python 3.12.10
torch 2.2.0+cu121
nequip 0.6.0
allegro 0.2.0

I tried adjusting this "neighbor 1.0 bin" in LAMMPS and r_max in my nequip yaml file.

Does anyone have any suggestions or explanations? Thank you!

[cw-tan]

I'm sorry but my only suggestion is that these are extremely outdated versions that will no longer be maintained, i.e. bugs in those versions will persist. Is there a specific reason why you wish to use those old versions? I'd strongly recommend just using the latest versions of nequip and allegro wherever possible, along with the up-to-date https://github.com/mir-group/pair_nequip_allegro