Source code for torch_lattice.nn.modules.activation

from __future__ import annotations

from torch import nn

from torch_lattice import SparseTensor
from torch_lattice.nn.utils import fapply

__all__ = [
    "GELU",
    "LeakyReLU",
    "ReLU",
    "SiLU",
    "Sigmoid",
    "Softplus",
    "Tanh",
]



[docs]
class ReLU(nn.ReLU):

[docs]
    def forward(self, input: SparseTensor) -> SparseTensor:
        return fapply(input, super().forward)





[docs]
class LeakyReLU(nn.LeakyReLU):

[docs]
    def forward(self, input: SparseTensor) -> SparseTensor:
        return fapply(input, super().forward)





[docs]
class SiLU(nn.SiLU):

[docs]
    def forward(self, input: SparseTensor) -> SparseTensor:
        return fapply(input, super().forward)





[docs]
class GELU(nn.GELU):

[docs]
    def forward(self, input: SparseTensor) -> SparseTensor:
        return fapply(input, super().forward)





[docs]
class Sigmoid(nn.Sigmoid):

[docs]
    def forward(self, input: SparseTensor) -> SparseTensor:
        return fapply(input, super().forward)





[docs]
class Tanh(nn.Tanh):

[docs]
    def forward(self, input: SparseTensor) -> SparseTensor:
        return fapply(input, super().forward)





[docs]
class Softplus(nn.Softplus):

[docs]
    def forward(self, input: SparseTensor) -> SparseTensor:
        return fapply(input, super().forward)