Activation Functions as Verified Operators

Implements element-wise activation functions (ReLU, Sigmoid, Tanh) as verified operators with kernel fusion support.

These are perfect examples of fusable operations because they:

1. Are element-wise (no cross-element dependencies)
2. Compose nicely: MatMul |> ReLU |> Softmax
3. Can be fused into preceding operations

Mathematical Definitions

ReLU: f(x) = max(0, x)

• Forward: y = max(0, x)
• Backward: dy/dx = 1 if x > 0, else 0

Sigmoid: f(x) = 1 / (1 + exp(-x))

• Forward: y = sigmoid(x)
• Backward: dy/dx = y * (1 - y)

Tanh: f(x) = tanh(x)

• Forward: y = tanh(x)
• Backward: dy/dx = 1 - y²

inductive Hesper.Op.Activation.ActivationType : Type

Activation function type

• Identity : ActivationType
• ReLU : ActivationType
• Sigmoid : ActivationType
• Gelu : ActivationType

def Hesper.Op.Activation.instInhabitedActivationType.default : ActivationType

Equations

• Hesper.Op.Activation.instInhabitedActivationType.default = Hesper.Op.Activation.ActivationType.Identity

instance Hesper.Op.Activation.instInhabitedActivationType : Inhabited ActivationType

Equations

• Hesper.Op.Activation.instInhabitedActivationType = { default := Hesper.Op.Activation.instInhabitedActivationType.default }

instance Hesper.Op.Activation.instReprActivationType : Repr ActivationType

Equations

• Hesper.Op.Activation.instReprActivationType = { reprPrec := Hesper.Op.Activation.instReprActivationType.repr }

def Hesper.Op.Activation.instReprActivationType.repr : ActivationType → Nat → Std.Format

Equations

• One or more equations did not get rendered due to their size.

instance Hesper.Op.Activation.instBEqActivationType : BEq ActivationType

Equations

• Hesper.Op.Activation.instBEqActivationType = { beq := Hesper.Op.Activation.instBEqActivationType.beq }

def Hesper.Op.Activation.instBEqActivationType.beq : ActivationType → ActivationType → Bool

Equations

• Hesper.Op.Activation.instBEqActivationType.beq x✝ y✝ = (x✝.ctorIdx == y✝.ctorIdx)

ReLU Activation

structure Hesper.Op.Activation.ReLUInput : Type

ReLU input: single tensor

• x : Core.TensorData

instance Hesper.Op.Activation.instInhabitedReLUInput : Inhabited ReLUInput

Equations

• Hesper.Op.Activation.instInhabitedReLUInput = { default := Hesper.Op.Activation.instInhabitedReLUInput.default }

def Hesper.Op.Activation.instInhabitedReLUInput.default : ReLUInput

Equations

• Hesper.Op.Activation.instInhabitedReLUInput.default = { x := default }

structure Hesper.Op.Activation.ReLUOutput : Type

ReLU output: single tensor

• y : Core.TensorData

def Hesper.Op.Activation.instInhabitedReLUOutput.default : ReLUOutput

Equations

• Hesper.Op.Activation.instInhabitedReLUOutput.default = { y := default }

instance Hesper.Op.Activation.instInhabitedReLUOutput : Inhabited ReLUOutput

Equations

• Hesper.Op.Activation.instInhabitedReLUOutput = { default := Hesper.Op.Activation.instInhabitedReLUOutput.default }

def Hesper.Op.Activation.cpuReLU (input : ReLUInput) : ReLUOutput

CPU ReLU forward pass: element-wise max(0, x)

Equations

• Hesper.Op.Activation.cpuReLU input = { y := { shape := input.x.shape, data := Array.map (fun (val : Float) => max 0.0 val) input.x.data } }

def Hesper.Op.Activation.cpuReLUBackward (input : ReLUInput) (grad_output : ReLUOutput) : ReLUInput

CPU ReLU backward pass: gradient is 1 if x > 0, else 0

Equations

• One or more equations did not get rendered due to their size.

GPU ReLU Kernel

def Hesper.Op.Activation.reluExp (x : WGSL.Exp (WGSL.WGSLType.scalar WGSL.ScalarType.f32)) : WGSL.Exp (WGSL.WGSLType.scalar WGSL.ScalarType.f32)

ReLU as a WGSL expression transformation. This is a pure function that can be fused into other kernels.

Equations

• Hesper.Op.Activation.reluExp x = x.max (Hesper.WGSL.Exp.litF32 0.0)

def Hesper.Op.Activation.gpuReLUKernel {_N : Nat} : WGSL.Kernel 256 1 1 (WGSL.Exp (WGSL.WGSLType.scalar WGSL.ScalarType.f32)) (WGSL.Exp (WGSL.WGSLType.scalar WGSL.ScalarType.f32))

GPU ReLU forward kernel. Element-wise operation that can be fused with other kernels.

Example usage:

let fused = matmul_kernel |> relu_kernel  -- Fuses matmul + relu

Equations

• Hesper.Op.Activation.gpuReLUKernel = Hesper.WGSL.mapK Hesper.Op.Activation.reluExp

def Hesper.Op.Activation.gpuReLUBackwardKernel : WGSL.Kernel 256 1 1 (WGSL.Exp (WGSL.WGSLType.scalar WGSL.ScalarType.f32) × WGSL.Exp (WGSL.WGSLType.scalar WGSL.ScalarType.f32)) (WGSL.Exp (WGSL.WGSLType.scalar WGSL.ScalarType.f32))

GPU ReLU backward kernel. Gradient: 1 if x > 0, else 0

Takes (input, grad_output) and returns grad_input. For ReLU: grad_input = grad_output if input > 0, else 0

Equations

• One or more equations did not get rendered due to their size.

VerifiedOpFusion Instance for ReLU

instance Hesper.Op.Activation.instVerifiedOpFusionOfNatNatReLUInputReLUOutputExpScalarF32 : Core.VerifiedOpFusion 256 1 1 ReLUInput ReLUOutput (WGSL.Exp (WGSL.WGSLType.scalar WGSL.ScalarType.f32)) (WGSL.Exp (WGSL.WGSLType.scalar WGSL.ScalarType.f32))

Equations

• One or more equations did not get rendered due to their size.

Helper Functions

def Hesper.Op.Activation.mkReLUInput (data : Core.TensorData) : ReLUInput

Create ReLU input from TensorData

Equations

• Hesper.Op.Activation.mkReLUInput data = { x := data }

def Hesper.Op.Activation.getReLUOutput (output : ReLUOutput) : Core.TensorData

Extract TensorData from ReLU output

Equations

• Hesper.Op.Activation.getReLUOutput output = output.y

Sigmoid Activation (TODO)

def Hesper.Op.Activation.sigmoidExp (x : WGSL.Exp (WGSL.WGSLType.scalar WGSL.ScalarType.f32)) : WGSL.Exp (WGSL.WGSLType.scalar WGSL.ScalarType.f32)

Sigmoid as WGSL expression: 1 / (1 + exp(-x))

Equations

• Hesper.Op.Activation.sigmoidExp x = (Hesper.WGSL.Exp.litF32 1.0).div ((Hesper.WGSL.Exp.litF32 1.0).add x.neg.exp)

Tanh Activation (TODO)

def Hesper.Op.Activation.tanhExp (x : WGSL.Exp (WGSL.WGSLType.scalar WGSL.ScalarType.f32)) : WGSL.Exp (WGSL.WGSLType.scalar WGSL.ScalarType.f32)

Tanh as WGSL expression

Equations

• Hesper.Op.Activation.tanhExp x = (x.exp.sub x.neg.exp).div (x.exp.add x.neg.exp)