Source file `activations.ml`

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
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
open Rune

(* Standard Activations *)

let relu = relu
let relu6 = relu6
let sigmoid = sigmoid
let tanh = tanh
let softmax = softmax

(* Modern Activations *)

let gelu = gelu
let silu = silu
let swish = silu (* Alias for silu *)
let mish = mish

(* Parametric Activations *)

let leaky_relu = leaky_relu
let elu = elu
let selu = selu

let prelu alpha x =
  (* max(0, x) + alpha * min(0, x) *)
  let zero = zeros_like x in
  add (maximum zero x) (mul alpha (minimum zero x))

(* Gated Linear Units (GLUs) *)

let glu x gate =
  (* x * sigmoid(gate) *)
  mul x (sigmoid gate)

let swiglu x =
  (* x * silu(x) *)
  mul x (silu x)

let geglu x gate =
  (* x * gelu(gate) *)
  mul x (gelu gate)

let reglu x gate =
  (* x * relu(gate) *)
  mul x (relu gate)

(* Other Activations *)

let softplus = softplus
let softsign = softsign
let hard_sigmoid = hard_sigmoid
let hard_tanh = hard_tanh

let hard_swish x =
  (* x * relu6(x + 3) / 6 - Not in Rune, but hard_silu is *)
  hard_silu x (* hard_silu is essentially the same as hard_swish *)