Multi Replica PDF

[APJansen]

NOTE: This is a copied and updated version of the PR #1782 which really should have been an issue.

Idea

The idea of this PR is to refactor the tensorflow model from taking a list of single-replica pdfs into taking a single multiple replica pdf, a single pdf whose output has an extra axis representing the replica. This is much faster on the GPU, see tests below.

The main ingredient to make this possible is a MultiDense layer, (see here) which is essentially just a dense layer where the weights have one extra dimension, with size the number of replicas. For the first layer, which takes x's as input, this is exactly it. For deeper layers, the input already has a replica axis, and so the right index of the input has to be multiplied by the corresponding axis of the weights.

Development Strategy

To integrate this into the code, many small changes are necessary.
To make it as simple as possible to review and test, I aim to make small, independent changes that ideally are beneficial, or at least not detrimental, on their own. Wherever it's sensible I'll first create a unit test that covers the changes I want to make, and make sure it still passes after, and wherever possible I'll try to have the outputs be identical up to numerical errors. I'll put all of these on their own branch and with their own PR (maybe I should create a special label for those PRs?).

Once those small changes are merged, the actual implementation should be easily managable to review.

I expect that as a final result you'll still want single replica pdf. I will add code that, once all computations are done, just splits the multi replica pdf into single ones, so the saving and any interaction with validphys will remain unchanged.

Performance

In the PR this was stated for the "sloppy" implementation, here I'll put the actual one once done.

Status

branch	finished	tested	merged	comments
refactor_xintegrator	X	unit	X
refactor_msr	X	unit	X
refactor_preprocessing	X	unit	X
refactor_rotations	X	unit	X
refactor_stopping	X	unit	X
multi-dense-logistics	X	X	X
replica-axis-first	X	X	X	ready for review (and done here because it makes the next one easier)
nn-layers-refactor	X	X		Currently layers are made with an outer loop over replicas and an inner loop over the depth in the network. It would be good to reverse this first, should be doable with identical results.
parallel-prefactor	X			WIP
multi-dense-layer	X

[RoyStegeman]

This has all been done. Thanks Aron!