File size: 848 Bytes
db6a3b7 a6bbecf db6a3b7 a6bbecf db6a3b7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 |
import torch
import torch.nn as nn
from . import SparseTensor
__all__ = ["SparseReLU", "SparseSiLU", "SparseGELU", "SparseActivation"]
class SparseReLU(nn.ReLU):
def forward(self, input: SparseTensor) -> SparseTensor:
return input.replace(super().forward(input.feats))
class SparseSiLU(nn.SiLU):
def forward(self, input: SparseTensor) -> SparseTensor:
return input.replace(super().forward(input.feats))
class SparseGELU(nn.GELU):
def forward(self, input: SparseTensor) -> SparseTensor:
return input.replace(super().forward(input.feats))
class SparseActivation(nn.Module):
def __init__(self, activation: nn.Module):
super().__init__()
self.activation = activation
def forward(self, input: SparseTensor) -> SparseTensor:
return input.replace(self.activation(input.feats))
|