jax empty element in the parameters of neural network

user_13604912 0
Viewed 46

I am working on the implementation of a very small neural network. My network is as follows:

init_random_params, predict = stax.serial(
Dense(1024), Relu,
Dense(1024), Relu,
Dense(10), LogSoftmax)

I have initialized the neural network as follows:

_, init_params = init_random_params(rng, (-1, 28 * 28))
params = init_params

when I print the parameters of neural network it is a tuple of size 6. In fact param[0] and param[2] and param[4] are non-empty and consists of weight and bias of the corresponding layers. However, the rest of elements, i.e., param[1] and param[3] and param[5] are simply empty.

I would like to understand the reason behind this because it makes it difficult for me to implement some training algorithms that I am interested.

autograd jax
1 Answers

Is there any particular reason you're using stax? Quoting from the stax module docstring:

You likely do not mean to import this module! Stax is intended as an example library only. There are a number of other much more fully-featured neural network libraries for JAX, including Flax from Google, and Haiku from DeepMind.

If you're having trouble creating neural networks with stax, you might try using an actively supported neural network library instead.

That said, I believe the reason there are no parameters in these entries is because ReLu and LogSoftmax do not require any parameters.

user_29378310
Related