core.models.equiformer_v2.heads.rank2

Copyright (c) Meta, Inc. and its affiliates.

This source code is licensed under the MIT license found in the LICENSE file in the root directory of this source tree.

Classes

Rank2Block	Output block for predicting rank-2 tensors (stress, dielectric tensor).
Rank2DecompositionEdgeBlock	Output block for predicting rank-2 tensors (stress, dielectric tensor, etc).
Rank2SymmetricTensorHead	A rank 2 symmetric tensor prediction head.

Module Contents

class core.models.equiformer_v2.heads.rank2.Rank2Block(emb_size: int, num_layers: int = 2, edge_level: bool = False, extensive: bool = False)

Bases: torch.nn.Module

Output block for predicting rank-2 tensors (stress, dielectric tensor). Applies outer product between edges and computes node-wise or edge-wise MLP.

Parameters:

• emb_size (int) – Size of edge embedding used to compute outer product

• num_layers (int) – Number of layers of the MLP

• edge_level (bool) – If true apply MLP at edge level before pooling, otherwise use MLP at nodes after pooling

• extensive (bool) – Whether to sum or average the outer products

edge_level

emb_size

extensive

scalar_nonlinearity

r2tensor_MLP

forward(edge_distance_vec, x_edge, edge_index, data)

Parameters:

• edge_distance_vec (torch.Tensor) – Tensor of shape (…, 3)

• x_edge (torch.Tensor) – Tensor of shape (…, emb_size)

• edge_index (torch.Tensor) – Tensor of shape (2, nEdges)

• data – LMDBDataset sample

class core.models.equiformer_v2.heads.rank2.Rank2DecompositionEdgeBlock(emb_size: int, num_layers: int = 2, edge_level: bool = False, extensive: bool = False)

Bases: torch.nn.Module

Output block for predicting rank-2 tensors (stress, dielectric tensor, etc). Decomposes a rank-2 symmetric tensor into irrep degree 0 and 2.

Parameters:

• emb_size (int) – Size of edge embedding used to compute outer product

• num_layers (int) – Number of layers of the MLP

• edge_level (bool) – If true apply MLP at edge level before pooling, otherwise use MLP at nodes after pooling

• extensive (bool) – Whether to sum or average the outer products

emb_size

edge_level

extensive

scalar_nonlinearity

scalar_MLP

irrep2_MLP

change_mat

forward(edge_distance_vec, x_edge, edge_index, data)

Parameters:

• edge_distance_vec (torch.Tensor) – Tensor of shape (…, 3)

• x_edge (torch.Tensor) – Tensor of shape (…, emb_size)

• edge_index (torch.Tensor) – Tensor of shape (2, nEdges)

• data – LMDBDataset sample

class core.models.equiformer_v2.heads.rank2.Rank2SymmetricTensorHead(backbone: fairchem.core.models.base.BackboneInterface, output_name: str = 'stress', decompose: bool = False, edge_level_mlp: bool = False, num_mlp_layers: int = 2, use_source_target_embedding: bool = False, extensive: bool = False, avg_num_nodes: int = 1.0, default_norm_type: str = 'layer_norm_sh')

Bases: torch.nn.Module, fairchem.core.models.base.HeadInterface

A rank 2 symmetric tensor prediction head.

ouput_name

name of output prediction property (ie, stress)

sphharm_norm

layer normalization for spherical harmonic edge weights

xedge_layer_norm

embedding layer norm

block

rank 2 equivariant symmetric tensor block

output_name

decompose

use_source_target_embedding

avg_num_nodes

sphharm_norm

xedge_layer_norm

forward(data: dict[str, torch.Tensor] | torch.Tensor, emb: dict[str, torch.Tensor]) → dict[str, torch.Tensor]

Parameters:

• data – data batch

• emb – dictionary with embedding object and graph data

Returns: dict of {output property name: predicted value}