structs for Neural Operators

[ayushinav]

fixes #16

[ayushinav]

I thought overloading Lux.setup and defining model(x, ps, st) would work without changing anything else but then there are other Lux functions that dispatch on a Lux.AbstractExplictLayer so maybe we can just always call model.ch everywhere?

[avik-pal]

I thought overloading Lux.setup

That is an invalid overload. It almost certainly won't do what you are expecting it to do.

I don't understand the rest of your question.

[ayushinav]

Instead of moving to structs, it is more convinient to make FNOs a compact layer too. I previously had a struct with the chain ch as its only parameter. This meant we call model.ch everytime, which did not seem very elegant to me, especially for someone using neural operators independently. I overloaded Lux.setup and model(x, ps, st) but there are other Lux functions which take a model and it would be somewhat hard to maintain if overload all the required functions.
Finally, moving FNO to compact layer works, and does not require any change in appplying the models.
To dispatch a function on DeepONet, we have CompactLuxLayer{:DeepONet} and similarly for FNO: CompactLuxLayer{:FourierNeuralOperator}
In future, we would therefore need to have all the models as compact layers or something similar if we want them to be dispatched in a similar way.
@avik-pal

That is an invalid overload. It almost certainly won't do what you are expecting it to do.

I don't understand the rest of your question.

In the structs setup, I had calling Lux.setup(rng, model) = Lux.setup(rng, model.ch) and similarly for model(x, ps, st)= model.ch(x, ps, st) but there are functions such as Lux.Experimental.single_train_step!(...) that take a Lux layer and it won't be easy to account for all such functions

[ayushinav]

5b6ce9f#diff-c2cb5b240d81999115bf4cfd0c46952645c08c0b620b37c91e51de853ccb6f51L18-L36

This was the overload I meant.

[avik-pal]

No that is not the correct way to define the structs discussed in that issue. Go through https://lux.csail.mit.edu/stable/manual/interface

[ayushinav]

I'll move around the docstring. Right now ig they do not correspond to right function calls.
Also, this might be a good time to convert OperatorKernel into struct too.

[avik-pal]

Yes let's do that.