Initializers

For that moment when built in initializers don’t satisfy your needs.

I found it weird that the flax library didn’t provide a “bounded” uniform initializer, so I took the code for the uniform one and modified it in a way that it allowed setting the bounds of the desired uniform initialization.

---

source

bounded_uniform

 bounded_uniform (minval=0.0, maxval=1.0, dtype=<class 'numpy.float64'>)

layer = nn.Dense(features=1000, kernel_init=bounded_uniform(minval=-10., maxval=-5.))
params = layer.init(random.PRNGKey(0), jnp.ones(shape=(1,32,32,3)))
plt.hist(params["params"]["kernel"].ravel())
plt.show()

2023-10-26 11:19:39.543311: E external/xla/xla/stream_executor/cuda/cuda_driver.cc:268] failed call to cuInit: CUDA_ERROR_NO_DEVICE: no CUDA-capable device is detected
No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

layer = nn.Dense(features=1000, kernel_init=bounded_uniform(minval=-10., maxval=25.))
params = layer.init(random.PRNGKey(0), jnp.ones(shape=(1,32,32,3)))
plt.hist(params["params"]["kernel"].ravel())
plt.show()

---

source

displaced_normal

 displaced_normal (mean=0.0, stddev=0.01, dtype=<class 'numpy.float64'>)

Builds an initializer that returns real normally-distributed random arrays.

	Type	Default	Details
mean	float	0.0	Mean of the distribution.
stddev	float	0.01	Standard deviation of the distribution.
dtype	type	float64	Desired DType of the resulting array.

layer = nn.Dense(features=1000, kernel_init=displaced_normal())
params = layer.init(random.PRNGKey(0), jnp.ones(shape=(1,32,32,3)))
plt.hist(params["params"]["kernel"].ravel())
plt.show()

layer = nn.Dense(features=1000, kernel_init=displaced_normal(mean=5., stddev=2.))
params = layer.init(random.PRNGKey(0), jnp.ones(shape=(1,32,32,3)))
plt.hist(params["params"]["kernel"].ravel())
plt.show()

---

source

freq_scales_init

 freq_scales_init (n_scales, fs, dtype=<class 'numpy.float64'>)

	Type	Default	Details
n_scales			Number of scales.
fs			Sampling frequency.
dtype	type	float64	Desired DType of the resulting array.

layer = nn.Dense(features=4, kernel_init=freq_scales_init(n_scales=4, fs=64))
params = layer.init(random.PRNGKey(0), jnp.ones(shape=(1,4)))
params

FrozenDict({
    params: {
        kernel: Array([24., 12.,  6.,  3.], dtype=float32),
        bias: Array([0., 0., 0., 0.], dtype=float32),
    },
})

---

source

k_array

 k_array (k, arr, dtype=<class 'numpy.float64'>)

	Type	Default	Details
k			Number of scales.
arr			Sampling frequency.
dtype	type	float64	Desired DType of the resulting array.

layer = nn.Dense(features=4, kernel_init=k_array(k=2, arr=jnp.array([24., 12., 6., 3.])))
params = layer.init(random.PRNGKey(0), jnp.ones(shape=(1,4)))
params

FrozenDict({
    params: {
        kernel: Array([0.08333334, 0.16666667, 0.33333334, 0.6666667 ], dtype=float32),
        bias: Array([0., 0., 0., 0.], dtype=float32),
    },
})

---

source

log_k_array

 log_k_array (k, arr, dtype=<class 'numpy.float64'>)

Initializer that generates the weights based on applying the log to a given array.

	Type	Default	Details
k			Number of scales.
arr			Sampling frequency.
dtype	type	float64	Desired DType of the resulting array.

layer = nn.Dense(features=4, kernel_init=log_k_array(k=1, arr=1/jnp.array([24., 12., 6., 3.])**2))
params = layer.init(random.PRNGKey(0), jnp.ones(shape=(1,4)))
params

FrozenDict({
    params: {
        kernel: Array([-6.3561077, -4.9698133, -3.583519 , -2.1972246], dtype=float32),
        bias: Array([0., 0., 0., 0.], dtype=float32),
    },
})

---

source

linspace

 linspace (start, stop, num, dtype=<class 'numpy.float64'>)

	Type	Default	Details
start
stop
num
dtype	type	float64	Desired DType of the resulting array.

layer = nn.Dense(features=4, kernel_init=linspace(start=0, stop=jnp.pi, num=4))
params = layer.init(random.PRNGKey(0), jnp.ones(shape=(1,4)))
params

FrozenDict({
    params: {
        kernel: Array([0.       , 0.7853982, 1.5707964, 2.3561945], dtype=float32),
        bias: Array([0., 0., 0., 0.], dtype=float32),
    },
})

---

source

equal_to

 equal_to (arr, dtype=<class 'numpy.float64'>)

	Type	Default	Details
arr
dtype	type	float64	Desired DType of the resulting array.

layer = nn.Dense(features=4, kernel_init=equal_to([1., 2., 3., 4.]))
params = layer.init(random.PRNGKey(0), jnp.ones(shape=(1,4)))
params

FrozenDict({
    params: {
        kernel: Array([1., 2., 3., 4.], dtype=float32),
        bias: Array([0., 0., 0., 0.], dtype=float32),
    },
})

---

source

mean

 mean (dtype=<class 'numpy.float64'>)

Builds an initializer that returns a kernel that calculates the mean of the interacting pixels.

	Type	Default	Details
dtype	type	float64	Desired DType of the resulting array.

layer = nn.Conv(features=1, kernel_size=(1,1), kernel_init=mean())
params = layer.init(random.PRNGKey(0), jnp.ones(shape=(1,32,32,3)))
plt.hist(params["params"]["kernel"].ravel())
plt.show()