How to break kernel_fn of serial into multiple kernel_fns of Dense and Activation

[FabianoVeglianti]

Hi, I've written a code that creates a MLP using serial and returns init_fn, apply_fn, kernel_fn, layers_list where layers_list is a list of dense and relu:
layers_list = [stax.Dense(n_list[0], Cw_sqrt, Cb_sqrt, parameterization=parameterization)]
for i in range(1, L+1):
layers_list += [stax.Relu(), stax.Dense(n_list[i], Cw_sqrt, Cb_sqrt, parameterization=parameterization)]
After that, I've written a function that returns a dictionary of 3 lists: one of apply_fns, one of init_fns and one of kernel_fns
def get_nn_layers_functions(layer_list):
layers_functions = {}
t = list(zip(*layer_list))
layers_functions["init_fns"] = t[0]
layers_functions["apply_fns"] = t[1]
layers_functions["kernel_fns"] = t[2]
return layers_functions
Finally I can either call the apply_fn of the serial object or call the apply_fn of each layer on the output of the apply_fn of the previous layer, something like this:
layers_functions = get_nn_layers_functions(layer_list)
_, params_layer1 = layers_functions["init_fns"][0](key[0], x1.shape)
z1 = layers_functions["apply_fns"][0](params_layer1, x1)
a1 = layers_functions["apply_fns"][1]((), z1)
_, params_layer2 = layers_functions["init_fns"][2](key[1], a1.shape)
z2 = layers_functions["apply_fns"][2](params_layer2, a1)

Can I do the same thing, so using the kernel_fns of the layers to reproduce the output I get with the kernel_fn of the serial object?

I tried, but I get an error.
You can find the code i produced on my gihutb, here link

[romanngg]

Thanks for the detailed repro! Here's a fix for the last cell in your colab

### Test Cell
...
kernel_0_ouptut = layers_functions["kernel_fns"][0](x1.reshape(1,16), None)
print(kernel_0_ouptut)

kernel_1_ouptut = layers_functions["kernel_fns"][1](kernel_0_ouptut)
print(kernel_1_ouptut.nngp)

The change lies in passing the whole Kernel dataclass as the input to the kernel_fn as opposed to Kernel.nngp which is just an array. When it's only an array, it misses the necessary metadata and actually interpreted as if you passed an input (akin to x1), and this raises an error when it's passed to the ReLU function, the infinite width limit of which requires it to be proceeded with a Gaussian linear layer. But when you pass the whole Kernel, the object contains the metadata (is_gaussian=True) about that.

[Also, on a minor note, it's slightly more efficient to pass (x, None) as opposed to (x, x) to the kernel function, where None means x2 == x1]

[romanngg]

Also, if you're interested at getting kernels at all layers of your network, check out this snippet which could be helpful #24 (comment)