per_residue_local_distance_difference_test

src/beignet/nn/functional/_per_residue_local_distance_difference_test.py

@@ -1,3 +1,5 @@
+# ruff: noqa: E501
+
 import torch
 import torch.nn.functional
 from torch import Tensor
@@ -13,17 +15,10 @@ def per_residue_local_distance_difference_test(input: Tensor) -> Tensor:
     -------
     output : Tensor
     """
-    probs = torch.nn.functional.softmax(input, dim=-1)
-
-    bins = input.shape[-1]
+    output = torch.nn.functional.softmax(input, dim=-1)
 
-    step = 1.0 / bins
+    step = 1.0 / input.shape[-1]
 
     bounds = torch.arange(0.5 * step, 1.0, step)
 
-    indexes = (1,) * len(probs.shape[:-1])
-    output = bounds.view(*indexes, *bounds.shape)
-    output = probs * output
-    output = torch.sum(output, dim=-1)
-
-    return output * 100
+    return torch.sum(output * torch.reshape(bounds, [*[1 for _ in range(len(output.shape[:-1]))], *bounds.shape]), dim=-1) * 100.0  # fmt: off