MolecularDiffusion.modules.layers.esen.nn.activation

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

GateActivation
S2Activation	Assume we only have one resolution
SeparableS2Activation

Module Contents

class MolecularDiffusion.modules.layers.esen.nn.activation.GateActivation(lmax: int, mmax: int, num_channels: int)

Bases: torch.nn.Module

forward(gating_scalars, input_tensors)

gating_scalars: shape [N, lmax * num_channels] input_tensors: shape [N, (lmax + 1) ** 2, num_channels]

gate_act

lmax

mmax

num_channels

scalar_act

class MolecularDiffusion.modules.layers.esen.nn.activation.S2Activation(lmax: int, mmax: int, SO3_grid)

Bases: torch.nn.Module

Assume we only have one resolution

forward(inputs)

SO3_grid

act

lmax

mmax

class MolecularDiffusion.modules.layers.esen.nn.activation.SeparableS2Activation(lmax: int, mmax: int, SO3_grid)

Bases: torch.nn.Module

forward(input_scalars, input_tensors)

lmax

mmax

s2_act

scalar_act