%install-location $cwd/swift-install
%install '.package(path: "$cwd/FastaiNotebook_02a_why_sqrt5")' FastaiNotebook_02a_why_sqrt5

//export
import Path
import TensorFlow

import FastaiNotebook_02a_why_sqrt5

var x = TF(randomNormal: [512, 1])
let a = TF(randomNormal: [512,512])

for i in 0..<100 { x = a • x }

(x.mean(),x.std())

var x = TF(randomNormal: [512, 1])
let a = TF(randomNormal: [512,512])

for i in 0..<100 {
    x = a • x
    if x.std().scalarized().isNaN {
        print(i)
        break
    }
}

var x = TF(randomNormal: [512, 1])
let a = TF(randomNormal: [512,512]) * 0.01

for i in 0..<100 { x = a • x }

(x.mean(),x.std())

var x = TF(randomNormal: [512, 1])
let a = TF(randomNormal: [512,512]) / sqrt(512)

for i in 0..<100 { x = a • x }

(mean: x.mean(), std: x.std())

1 / sqrt(512)

var x = TF(randomNormal: [512, 1])
(mean: x.mean(), std: x.std())

var mean = Float()
var sqr = Float()
for i in 0..<100 {
    let x = TF(randomNormal: [512, 1])
    let a = TF(randomNormal: [512, 512])
    let y = a • x
    mean += y.mean().scalarized()
    sqr  += pow(y, 2).mean().scalarized()
}
(mean/100, sqr/100)

var mean = Float()
var sqr = Float()
for i in 0..<10000 {
    let x = TF(randomNormal: [])
    let a = TF(randomNormal: [])
    let y = a * x
    mean += y.scalarized()
    sqr  += pow(y, 2).scalarized()
}
(mean/10000,sqrt(sqr/10000))

var mean = Float()
var sqr = Float()
for i in 0..<10000 {
    let x = TF(randomNormal: [])
    let a = TF(randomNormal: [])
    var y = (a*x).scalarized()
    y = y < 0 ? 0 : y
    mean += y
    sqr  += pow(y, 2)
}
(mean: mean/10000, sqrt: sqr/10000)

var mean = Float()
var sqr = Float()
for i in 0..<100 {
    let x = TF(randomNormal: [512, 1])
    let a = TF(randomNormal: [512, 512])
    var y = a • x
    y = max(y, TF(zeros: y.shape))
    mean += y.mean().scalarized()
    sqr  += pow(y, 2).mean().scalarized()
}
(mean: mean/100, sqrt: sqr/100)

var mean = Float()
var sqr = Float()
for i in 0..<100 {
    let x = TF(randomNormal: [512, 1])
    let a = TF(randomNormal: [512, 512]) * sqrt(2/512)
    var y = a • x
    y = max(y, TF(zeros: y.shape))
    mean += y.mean().scalarized()
    sqr  += pow(y, 2).mean().scalarized()
}
(mean: mean/100, sqrt: sqr/100)