xnn.gnn.featurizers.radial.IdentityDistanceTransform#

class xnn.gnn.featurizers.radial.IdentityDistanceTransform(*args, **kwargs)[source]#

Bases: Module

The trivial distance transform: distances pass through unchanged.

Keeps the edge-embedding code path uniform (and TorchScript-friendly) for models whose radial basis acts on the raw interatomic distance.

forward(r, atomic_numbers, edge_index)[source]#

Return the distances unchanged.

Parameters:

  • r (Tensor) – Interatomic distances, shape (E,).

  • atomic_numbers (Tensor) – Per-node atomic numbers, shape (N,) (unused).

  • edge_index (Tensor) – Edge index of shape (2, E) (unused).

Returns:

r, unchanged.

Return type:

Tensor