Switch objective function to be the first argument in Training API

[avik-pal]

This would allow us to write the objective function in block form like

Training.single_train_step!(backend, data, train_state) do model, ps, st, data
    ...
end

Recently posted as a gripe on slack, it is a fair ask, but there are no promises when we are doing a v2 release. But whenever we do it, I will gladly make the switch.