chop.nn

chop.nn.functional

chop.nn.functional.softermax.softermax(input: Tensor, dim: int) → Tensor

Softermax implementation, according to “Softermax: Hardware/Software Co-Design of an Efficient Softmax for Transformers” paper (https://arxiv.org/abs/2103.09301).

Parameters:

input (Tensor) – Input tensor

Returns:

Output tensor

Return type:

Tensor

chop.nn.functional.splitter.splitter(x)

chop.nn.modules