Closed murphyk closed 1 year ago
Following https://uvadlc-notebooks.readthedocs.io/en/latest/tutorial_notebooks/JAX/tutorial3/Activation_Functions.html#Setup, we can create our own kernel_init and bias_init functions, which sample from TruncGauss(0,a) and Gauss(0,0.0001) for weights and biases, where a=1/sqrt{Fin). This is the same as default flax init, but gives us explicit access to Sigma0, rather than having to hack the source code. We can use this for setting Upisilon_0 to be diagonal (but not spherical).
kernel_init
bias_init
TruncGauss(0,a)
Gauss(0,0.0001)
a=1/sqrt{Fin)
Following https://uvadlc-notebooks.readthedocs.io/en/latest/tutorial_notebooks/JAX/tutorial3/Activation_Functions.html#Setup, we can create our own
kernel_initandbias_initfunctions, which sample fromTruncGauss(0,a)andGauss(0,0.0001)for weights and biases, wherea=1/sqrt{Fin). This is the same as default flax init, but gives us explicit access to Sigma0, rather than having to hack the source code. We can use this for setting Upisilon_0 to be diagonal (but not spherical).