Incorrectly pytree recognition by KFAC optimizer

[Uernd]

First of all, thanks for the contribution of KFAC team !

While using KFAC optimizer to optimize an ANN, I noticed that the KFAC optimizer seems have some trouble to understand the structure of parameter tree if the parameter is used more than once while constructing the neural network.

If the original ANN denoted as f(params, inputs), then if we simply use a modified ANN as F(params, inputs) = f(params, inputs) + f(params, inputs)，the program will throw an error. I have tried functools.partial to fix the parameters, but it seems the program will get stuck somehow. If I use vmap, some of the parameters would be labelled as 'orphan' and by experiments, this would affect the optimization process.

I wonder if there is already some methods to avoid these issues? Would you consider update the optimizer to fix this bug?

Thanks again for the well-designed optimizer !