Activation modules

Shape-preserving activation specifications packaged for module composition.

def Spec.Module.relu {α : Type} [TorchLean.Storage α] [Zero α] [Max α] [BEq α] (s : Shape) : Module α s s

ReLU as a shape-preserving module.

def Spec.Module.sigmoid {α : Type} [TorchLean.Storage α] [Context α] (s : Shape) : Module α s s

Sigmoid as a shape-preserving module.

def Spec.Module.tanh {α : Type} [TorchLean.Storage α] [Context α] (s : Shape) : Module α s s

Hyperbolic tangent as a shape-preserving module.

def Spec.Module.softmax {α : Type} [TorchLean.Storage α] [Context α] (s : Shape) (axis : ℕ) [Shape.AxisInBounds axis s] : Module α s s

Softmax along an explicitly selected tensor dimension.

In PyTorch terms: torch.softmax(x, dim=axis).