Training nnx NN models without pre-defining number of layers/neurons

[yigitcancomlek]

Hello!

I am trying to convert my code in linen to nnx. My previous linen code was able to handle any given architecture for a simple multi-layer neural network through the call function. However, with nnx, it looks like I need to pre-define my layers in the init to achieve this, which is a bit inconvenient, considering if I want to do hyperparameter optimization or just simply try different architectures for my data. All the examples in the flax nnx version website also predefines architecture before training the model. Therefore, I was wondering if the developers could kindly provide a simple example of how to handle any architecture without predefining it in the class, that would be great! If this is not possible, assuming I understood this Q&A correctly, please let me know as well! Thanks a lot in advance!