model/nn/nonlinearities.py

# This file is a part of the `nequip` package. Please see LICENSE and README at the root for information on using it.
import torch

import math

# Technically, use of this module should probably be guarded by conditional_torchscript_jit
# But its use as a drop-in replacement for functions like torch.nn.functional.silu makes that
# difficult, so, given the rarety of its use, we have just removed @torch.jit.script


def shifted_softplus(x):
    return torch.nn.functional.softplus(x) - math.log(2.0)


class ShiftedSoftplus(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        return shifted_softplus(x)