Match Flux's support of arbitrary loss functions

[Vilin97]

In Flux/Zygote I can use any loss function, whereas SimpleChains only allow absolute loss, squared loss and cross-entropy loss (am I perhaps wrong?). What is the reason that an arbitrary loss cannot be used? I would want something like

using SimpleChains
y = rand(Float32, 3, 2^8)
s = SimpleChain(
  static(3),
  TurboDense{true}(softsign, 100),
  TurboDense{true}(softsign, 3)
);
my_loss_fun(arg) = sum(i -> (arg[i] - y[i])^2, eachindex(y))/sum(abs2, y) 
my_loss = SimpleChains.Loss(my_loss_fun, y) # proposed syntax
train_loss = SimpleChains.add_loss(s, my_loss) # proposed syntax

Is it possible to do something like this? What would it take to use my own loss with SimpleChains?

[chriselrod]

It'd probably be possible.
I don't have time to work on SimpleChains at the moment, but would be happy to discuss ideas.

Simple cases where lossan be computed elementwise could be calculated efficiently with ForwardDiff.
SimpleChains supports arbitrary activation functions in this way.

My long term plan is to get LoopModels working (ideally for both CPU and GPU targets), and get it playing well with Enzyme, so that we can cover generic code.
It'll be some time before we get that far.

But the current reality is that all our AD systems are bad, and SimpleChains does well on benchmarks thanks to not using them.
I just wrote all the rules I wanted.