MolecularDiffusion.modules.layers.esen.nn.so3_layers¶

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¶

Module Contents¶

class MolecularDiffusion.modules.layers.esen.nn.so3_layers.SO3_Linear(in_features: int, out_features: int, lmax: int)¶

Bases: torch.nn.Module

Use torch.einsum to prevent slicing and concatenation
Need to specify some behaviors in no_weight_decay and weight initialization.

forward(input_embedding)¶

bias¶

in_features¶

lmax¶

out_features¶

weight¶