File size: 1,976 Bytes
3e02ab8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 | from ._tp_scatter_base import TensorProductScatter
class OpenEquivarianceTensorProductScatter(TensorProductScatter):
_nequip_custom_ops_libs = ("openequivariance",)
def __init__(
self,
feature_irreps_in,
irreps_edge_attr,
irreps_mid,
instructions,
use_opaque: bool,
) -> None:
super().__init__(
feature_irreps_in=feature_irreps_in,
irreps_edge_attr=irreps_edge_attr,
irreps_mid=irreps_mid,
instructions=instructions,
)
# ^ we ensure that the base class keeps around a `self.tp` that carries its own set of persistent buffers
# even though `self.tp` is not used, having its (persistent) buffers always around ensures state dict compatibility when adding on or removing this subclass module
# === OEQ ===
# we do lazy imports of oeq to allow `nequip-package` to pick this file up even if oeq is not installed
# since `nequip-package` ignores files if it errors on loading the file
from openequivariance import (
TensorProductConv,
TPProblem,
torch_to_oeq_dtype,
)
tpp = TPProblem(
feature_irreps_in,
irreps_edge_attr,
irreps_mid,
instructions,
irrep_dtype=torch_to_oeq_dtype(self.model_dtype),
weight_dtype=torch_to_oeq_dtype(self.model_dtype),
shared_weights=False,
internal_weights=False,
)
self.tp_conv = TensorProductConv(
tpp, torch_op=True, deterministic=False, use_opaque=use_opaque
)
def forward(self, x, edge_attr, edge_weight, edge_dst, edge_src):
# explicit cast to account for AMP
return self.tp_conv(
x.to(self.model_dtype),
edge_attr.to(self.model_dtype),
edge_weight.to(self.model_dtype),
edge_dst,
edge_src,
)
|