Trax Extra Activation Functions

Extra activation functions for Trax

pip install trax-extra-activation

How to use

You can use activation functions in the same way as that ones implemented in Trax.

from trax_extra_activation import TanhExp
model = tl.Serial(
    tl.Embedding(vocab_size=8192, d_feature=256),
    tl.Mean(axis=1),
    tl.Dense(2),
    TanhExp()
)

Activations

• HardShrink

  HardShrink(lmbda: float = 0.5) -> PureLayer
  

  • PyTorch Document
• HardSwish

  HardSwish() -> PureLayer
  

  • PyTorch Document
• LogSigmoid

  LogSigmoid() -> PureLayer
  

  • PyTorch Document
• Relu6

  Relu6() -> PureLayer
  

  • PyTorch Doument
• Celu

  Celu(a: float = 1.0) -> PureLayer
  

  • PyTorch Document
• Silu

  Silu() -> PureLayer
  

  • PyTorch Document
• Mish

  Mish() -> PureLayer
  

  • PyTorch Document
• SoftShrink

  SoftShrink(lmbda: float = 0.5) -> PureLayer
  

  • PyTorch Document
• SoftSign

  SoftSign() -> PureLayer
  

  • PyTorch Document
• TanhShrink

  TanhShrink() -> PureLayer
  

  • PyTorch Document
• TanhExp

  TanhExp() -> PureLayer
  

  • Paper