diff --git a/n3fit/src/n3fit/backends/keras_backend/MetaModel.py b/n3fit/src/n3fit/backends/keras_backend/MetaModel.py
index 9ed23b0823..dfefd38401 100644
--- a/n3fit/src/n3fit/backends/keras_backend/MetaModel.py
+++ b/n3fit/src/n3fit/backends/keras_backend/MetaModel.py
@@ -396,8 +396,7 @@ def split_replicas(self):
         if self.single_replica_generator is None:
             raise ValueError("Trying to generate single replica models with no generator set.")
         replicas = []
-        num_replicas = self.output.shape[-1]
-        for i_replica in range(num_replicas):
+        for i_replica in range(self.num_replicas):
             replica = self.single_replica_generator()
             replica.set_replica_weights(self.get_replica_weights(i_replica))
 
@@ -410,14 +409,17 @@ def split_replicas(self):
 
         return replicas
 
+    @property
+    def num_replicas(self):
+        return self.output.shape[1]
+
     def load_identical_replicas(self, model_file):
         """
         From a single replica model, load the same weights into all replicas.
         """
         weights = self._format_weights_from_file(model_file)
 
-        num_replicas = self.output.shape[-1]
-        for i_replica in range(num_replicas):
+        for i_replica in range(self.num_replicas):
             self.set_replica_weights(weights, i_replica)
 
     def _format_weights_from_file(self, model_file):
diff --git a/n3fit/src/n3fit/backends/keras_backend/operations.py b/n3fit/src/n3fit/backends/keras_backend/operations.py
index 4ec1ecea4f..f965152f42 100644
--- a/n3fit/src/n3fit/backends/keras_backend/operations.py
+++ b/n3fit/src/n3fit/backends/keras_backend/operations.py
@@ -261,7 +261,7 @@ def pdf_masked_convolution(raw_pdf, basis_mask):
     Parameters
     ----------
         pdf: tf.tensor
-            rank 4 (batchsize, xgrid, flavours, replicas)
+            rank 4 (batchsize, replicas, xgrid, flavours)
         basis_mask: tf.tensor
             rank  2 tensor (flavours, flavours)
             mask to apply to the pdf convolution
@@ -269,18 +269,18 @@ def pdf_masked_convolution(raw_pdf, basis_mask):
     Return
     ------
         pdf_x_pdf: tf.tensor
-            rank3 (len(mask_true), xgrid, xgrid, replicas)
+            rank3 (replicas, len(mask_true), xgrid, xgrid)
     """
-    if raw_pdf.shape[-1] == 1:  # only one replica!
-        pdf = tf.squeeze(raw_pdf, axis=(0, -1))
+    if raw_pdf.shape[1] == 1:  # only one replica!
+        pdf = tf.squeeze(raw_pdf, axis=(0, 1))
         luminosity = tensor_product(pdf, pdf, axes=0)
         lumi_tmp = K.permute_dimensions(luminosity, (3, 1, 2, 0))
-        pdf_x_pdf = batchit(boolean_mask(lumi_tmp, basis_mask), -1)
+        pdf_x_pdf = batchit(boolean_mask(lumi_tmp, basis_mask), 0)
     else:
         pdf = tf.squeeze(raw_pdf, axis=0)  # remove the batchsize
-        luminosity = tf.einsum('air,bjr->jibar', pdf, pdf)
+        luminosity = tf.einsum('rai,rbj->rjiba', pdf, pdf)
         # (xgrid, flavour, xgrid, flavour)
-        pdf_x_pdf = boolean_mask(luminosity, basis_mask)
+        pdf_x_pdf = boolean_mask(luminosity, basis_mask, axis=1)
     return pdf_x_pdf
 
 
diff --git a/n3fit/src/n3fit/hyper_optimization/penalties.py b/n3fit/src/n3fit/hyper_optimization/penalties.py
index 414758ce38..77a156cd3a 100644
--- a/n3fit/src/n3fit/hyper_optimization/penalties.py
+++ b/n3fit/src/n3fit/hyper_optimization/penalties.py
@@ -53,17 +53,18 @@ def saturation(pdf_model=None, n=100, min_x=1e-6, max_x=1e-4, flavors=None, **_k
     if flavors is None:
         flavors = [1, 2]
     x = np.logspace(np.log10(min_x), np.log10(max_x), n)
-    x = np.expand_dims(x, axis=[0, -1])
     extra_loss = 0.0
 
-    y = pdf_model.predict({"pdf_input": x})
-    xpdf = y[0, :, flavors]
+    x_input = np.expand_dims(x, axis=[0, -1])
+    y = pdf_model.predict({"pdf_input": x_input})
+    xpdf = y[0, :, :, flavors]  # this is now of shape (flavors, replicas, xgrid)
 
-    delta_logx = np.diff(np.log10(x), axis=1)
-    delta_xpdf = np.diff(xpdf, axis=1)
+    x = np.expand_dims(x, axis=[0, 1])
+    delta_logx = np.diff(np.log10(x), axis=2)
+    delta_xpdf = np.diff(xpdf, axis=2)
     slope = delta_xpdf / delta_logx
 
-    pen = abs(np.mean(slope, axis=1)) + np.std(slope, axis=1)
+    pen = abs(np.mean(slope, axis=2)) + np.std(slope, axis=2)
 
     # sum over flavors
     # Add a small offset to avoid ZeroDivisionError
diff --git a/n3fit/src/n3fit/layers/DIS.py b/n3fit/src/n3fit/layers/DIS.py
index 70dec828f6..0aeb2d00e6 100644
--- a/n3fit/src/n3fit/layers/DIS.py
+++ b/n3fit/src/n3fit/layers/DIS.py
@@ -7,29 +7,31 @@
 """
 
 import numpy as np
-from .observable import Observable
+
 from n3fit.backends import operations as op
 
+from .observable import Observable
+
 
 class DIS(Observable):
     """
-        The DIS class receives a list of active flavours and a fktable
-        and prepares a layer that performs the convolution of said fktable with
-        the incoming pdf.
+    The DIS class receives a list of active flavours and a fktable
+    and prepares a layer that performs the convolution of said fktable with
+    the incoming pdf.
 
-        The fktable is expected to be rank 3 (ndata, xgrid, flavours)
-        while the input pdf is rank 4 where the first dimension is the batch dimension
-        and the last dimension the number of replicas being fitted (1, xgrid, flavours, replicas)
+    The fktable is expected to be rank 3 (ndata, xgrid, flavours)
+    while the input pdf is rank 4 where the first dimension is the batch dimension
+    and the last dimension the number of replicas being fitted (1, replicas, xgrid, flavours)
     """
 
     def gen_mask(self, basis):
         """
-            Receives a list of active flavours and generates a boolean mask tensor
+        Receives a list of active flavours and generates a boolean mask tensor
 
-            Parameters
-            ----------
-                basis: list(int)
-                    list of active flavours
+        Parameters
+        ----------
+            basis: list(int)
+                list of active flavours
         """
         if basis is None:
             self.basis = np.ones(self.nfl, dtype=bool)
@@ -41,21 +43,21 @@ def gen_mask(self, basis):
 
     def call(self, pdf):
         """
-            This function perform the fktable \otimes pdf convolution.
+        This function perform the fktable \otimes pdf convolution.
 
-            First pass the input PDF through a mask to remove the unactive flavors,
-            then a tensor_product between the PDF and each fktable is performed
-            finally the defined operation is applied to all the results
+        First pass the input PDF through a mask to remove the unactive flavors,
+        then a tensor_product between the PDF and each fktable is performed
+        finally the defined operation is applied to all the results
 
-            Parameters
-            ----------
-                pdf:  backend tensor
-                    rank 4 tensor (batch_size, xgrid, flavours, replicas)
+        Parameters
+        ----------
+            pdf:  backend tensor
+                rank 4 tensor (batch_size, replicas, xgrid, flavours)
 
-            Returns
-            -------
-                result: backend tensor
-                    rank 3 tensor (batchsize, replicas, ndata)
+        Returns
+        -------
+            result: backend tensor
+                rank 3 tensor (batchsize, replicas, ndata)
         """
         # DIS never needs splitting
         if self.splitting is not None:
@@ -65,13 +67,13 @@ def call(self, pdf):
         # Separate the two possible paths this layer can take
         if self.many_masks:
             for mask, fktable in zip(self.all_masks, self.fktables):
-                pdf_masked = op.boolean_mask(pdf, mask, axis=2)
-                res = op.tensor_product(pdf_masked, fktable, axes=[(1, 2), (2, 1)])
+                pdf_masked = op.boolean_mask(pdf, mask, axis=3)
+                res = op.tensor_product(pdf_masked, fktable, axes=[(2, 3), (2, 1)])
                 results.append(res)
         else:
-            pdf_masked = op.boolean_mask(pdf, self.all_masks[0], axis=2)
+            pdf_masked = op.boolean_mask(pdf, self.all_masks[0], axis=3)
             for fktable in self.fktables:
-                res = op.tensor_product(pdf_masked, fktable, axes=[(1, 2), (2, 1)])
+                res = op.tensor_product(pdf_masked, fktable, axes=[(2, 3), (2, 1)])
                 results.append(res)
 
         return self.operation(results)
diff --git a/n3fit/src/n3fit/layers/DY.py b/n3fit/src/n3fit/layers/DY.py
index dc5fe858da..3b42a394bd 100644
--- a/n3fit/src/n3fit/layers/DY.py
+++ b/n3fit/src/n3fit/layers/DY.py
@@ -1,7 +1,9 @@
 import numpy as np
-from .observable import Observable
+
 from n3fit.backends import operations as op
 
+from .observable import Observable
+
 
 class DY(Observable):
     """
@@ -30,7 +32,7 @@ def call(self, pdf_raw):
         Parameters
         ----------
             pdf_in: tensor
-                rank 4 tensor (batchsize, xgrid, flavours, replicas)
+                rank 4 tensor (batchsize, replicas, xgrid, flavours)
 
         Returns
         -------
@@ -43,20 +45,20 @@ def call(self, pdf_raw):
         results = []
         if self.many_masks:
             if self.splitting:
-                splitted_pdf = op.split(pdf_raw, self.splitting, axis=1)
+                splitted_pdf = op.split(pdf_raw, self.splitting, axis=2)
                 for mask, pdf, fk in zip(self.all_masks, splitted_pdf, self.fktables):
                     pdf_x_pdf = op.pdf_masked_convolution(pdf, mask)
-                    res = op.tensor_product(fk, pdf_x_pdf, axes=3)
+                    res = op.tensor_product(fk, pdf_x_pdf, axes=[(1, 2, 3), (1, 2, 3)])
                     results.append(res)
             else:
                 for mask, fk in zip(self.all_masks, self.fktables):
                     pdf_x_pdf = op.pdf_masked_convolution(pdf_raw, mask)
-                    res = op.tensor_product(fk, pdf_x_pdf, axes=3)
+                    res = op.tensor_product(fk, pdf_x_pdf, axes=[(1, 2, 3), (1, 2, 3)])
                     results.append(res)
         else:
             pdf_x_pdf = op.pdf_masked_convolution(pdf_raw, self.all_masks[0])
             for fk in self.fktables:
-                res = op.tensor_product(fk, pdf_x_pdf, axes=3)
+                res = op.tensor_product(fk, pdf_x_pdf, axes=[(1, 2, 3), (1, 2, 3)])
                 results.append(res)
 
         # the masked convolution removes the batch dimension
diff --git a/n3fit/src/n3fit/layers/msr_normalization.py b/n3fit/src/n3fit/layers/msr_normalization.py
index 0711eee86d..0755cad0b4 100644
--- a/n3fit/src/n3fit/layers/msr_normalization.py
+++ b/n3fit/src/n3fit/layers/msr_normalization.py
@@ -11,38 +11,14 @@
 from n3fit.backends import MetaLayer
 from n3fit.backends import operations as op
 
-IDX = {
-    'photon': 0,
-    'sigma': 1,
-    'g': 2,
-    'v': 3,
-    'v3': 4,
-    'v8': 5,
-    'v15': 6,
-    'v24': 7,
-    'v35': 8,
-}
+IDX = {'photon': 0, 'sigma': 1, 'g': 2, 'v': 3, 'v3': 4, 'v8': 5, 'v15': 6, 'v24': 7, 'v35': 8}
 MSR_COMPONENTS = ['g']
 MSR_DENOMINATORS = {'g': 'g'}
 # The VSR normalization factor of component f is given by
 # VSR_CONSTANTS[f] / VSR_DENOMINATORS[f]
 VSR_COMPONENTS = ['v', 'v35', 'v24', 'v3', 'v8', 'v15']
-VSR_CONSTANTS = {
-    'v': 3.0,
-    'v35': 3.0,
-    'v24': 3.0,
-    'v3': 1.0,
-    'v8': 3.0,
-    'v15': 3.0,
-}
-VSR_DENOMINATORS = {
-    'v': 'v',
-    'v35': 'v',
-    'v24': 'v',
-    'v3': 'v3',
-    'v8': 'v8',
-    'v15': 'v15',
-}
+VSR_CONSTANTS = {'v': 3.0, 'v35': 3.0, 'v24': 3.0, 'v3': 1.0, 'v8': 3.0, 'v15': 3.0}
+VSR_DENOMINATORS = {'v': 'v', 'v35': 'v', 'v24': 'v', 'v3': 'v3', 'v8': 'v8', 'v15': 'v15'}
 
 
 class MSR_Normalization(MetaLayer):
@@ -93,18 +69,20 @@ def call(self, pdf_integrated, photon_integral):
 
         Parameters
         ----------
-        pdf_integrated: (Tensor(1, 14, replicas))
+        pdf_integrated: (Tensor(1, replicas, 14))
             the integrated PDF
-        photon_integral: (Tensor(1, 1, replicas))
+        photon_integral: (Tensor(1, replicas, 1))
             the integrated photon PDF
 
         Returns
         -------
-        normalization_factor: Tensor(14, replicas)
+        normalization_factor: Tensor(replicas, 1, 14)
             The normalization factors per flavour.
         """
-        y = pdf_integrated[0]  # get rid of the batch dimension
-        photon_integral = photon_integral[0]  # get rid of the batch dimension
+        # get rid of batch dimension and put replicas last
+        reshape = lambda x: op.transpose(x[0])
+        y = reshape(pdf_integrated)
+        photon_integral = reshape(photon_integral)
         numerators = []
 
         if self._msr_enabled:
@@ -122,4 +100,4 @@ def call(self, pdf_integrated, photon_integral):
             numerators / divisors, indices=self.indices, output_shape=y.shape
         )
 
-        return norm_constants
+        return op.batchit(op.transpose(norm_constants), batch_dimension=1)
diff --git a/n3fit/src/n3fit/layers/rotations.py b/n3fit/src/n3fit/layers/rotations.py
index 4d9e4eaf8f..4ff9272682 100644
--- a/n3fit/src/n3fit/layers/rotations.py
+++ b/n3fit/src/n3fit/layers/rotations.py
@@ -22,7 +22,7 @@ class Rotation(MetaLayer):
             rotation_axis of input to be rotated
     """
 
-    def __init__(self, rotation_matrix, rotation_axis=2, **kwargs):
+    def __init__(self, rotation_matrix, rotation_axis=3, **kwargs):
         self.rotation_matrix = op.numpy_to_tensor(rotation_matrix)
         self.rotation_axis = rotation_axis
         super().__init__(**kwargs)
@@ -47,12 +47,7 @@ class FlavourToEvolution(Rotation):
     the evolution basis.
     """
 
-    def __init__(
-        self,
-        flav_info,
-        fitbasis,
-        **kwargs,
-    ):
+    def __init__(self, flav_info, fitbasis, **kwargs):
         rotation_matrix = pdfbases.fitbasis_to_NN31IC(flav_info, fitbasis)
         super().__init__(rotation_matrix, **kwargs)
 
@@ -111,7 +106,7 @@ def __init__(self, photons, **kwargs):
         super().__init__(**kwargs)
 
     def register_photon(self, xgrid):
-        """Compute the photon array of shape (1, xgrid, 1, replicas) and set the layer to be rebuilt"""
+        """Compute the photon array of shape (1, replicas, xgrid, 1) and set the layer to be rebuilt"""
         if self._photons_generator:
             self._pdf_ph = self._photons_generator(xgrid)
             self.built = False
@@ -119,7 +114,7 @@ def register_photon(self, xgrid):
     def call(self, pdfs):
         if self._pdf_ph is None:
             return pdfs
-        return op.concatenate([self._pdf_ph, pdfs[:, :, 1:]], axis=2)
+        return op.concatenate([self._pdf_ph, pdfs[:, :, :, 1:]], axis=3)
 
 
 class ObsRotation(MetaLayer):
diff --git a/n3fit/src/n3fit/layers/x_operations.py b/n3fit/src/n3fit/layers/x_operations.py
index 1d22b78cca..a55b56d2ba 100644
--- a/n3fit/src/n3fit/layers/x_operations.py
+++ b/n3fit/src/n3fit/layers/x_operations.py
@@ -65,11 +65,11 @@ class xIntegrator(MetaLayer):
     ----------
         grid_weights: np.array
             weights of the grid
-        x_axis: int (default=1)
+        x_axis: int (default=2)
             axis of the input tensor that corresponds to the x-grid
     """
 
-    def __init__(self, grid_weights, x_axis=1, **kwargs):
+    def __init__(self, grid_weights, x_axis=2, **kwargs):
         self.x_axis = x_axis
         self.grid_weights = op.flatten(op.numpy_to_tensor(grid_weights))
         super().__init__(**kwargs)
diff --git a/n3fit/src/n3fit/model_gen.py b/n3fit/src/n3fit/model_gen.py
index 87e66753d9..a700a10fa9 100644
--- a/n3fit/src/n3fit/model_gen.py
+++ b/n3fit/src/n3fit/model_gen.py
@@ -78,7 +78,7 @@ def _generate_experimental_layer(self, pdf):
             splitting_layer = op.as_layer(
                 op.split,
                 op_args=[self.dataset_xsizes],
-                op_kwargs={"axis": 1},
+                op_kwargs={"axis": 2},
                 name=f"{self.name}_split",
             )
             sp_pdf = splitting_layer(pdf)
@@ -578,7 +578,7 @@ def pdfNN_layer_generator(
     Returns
     -------
        pdf_model: n3fit.backends.MetaModel
-            a model f(x) = y where x is a tensor (1, xgrid, 1) and y a tensor (1, xgrid, out, num_replicas)
+            a model f(x) = y where x is a tensor (1, xgrid, 1) and y a tensor (1, replicas, xgrid, out)
     """
     # Parse the input configuration
     if seed is None:
@@ -700,13 +700,13 @@ def pdfNN_layer_generator(
 
     # Apply NN layers for all replicas to a given input grid
     def neural_network_replicas(x, postfix=""):
-        NNs_x = Lambda(lambda nns: op.stack(nns, axis=-1), name=f"NNs{postfix}")(
+        NNs_x = Lambda(lambda nns: op.stack(nns, axis=1), name=f"NNs{postfix}")(
             [nn(x) for nn in nn_replicas]
         )
 
         if subtract_one:
             x_eq_1_processed = process_input(layer_x_eq_1)
-            NNs_x_1 = Lambda(lambda nns: op.stack(nns, axis=-1), name=f"NNs{postfix}_x_1")(
+            NNs_x_1 = Lambda(lambda nns: op.stack(nns, axis=1), name=f"NNs{postfix}_x_1")(
                 [nn(x_eq_1_processed) for nn in nn_replicas]
             )
             NNs_x = subtract_one_layer([NNs_x, NNs_x_1])
@@ -715,7 +715,7 @@ def neural_network_replicas(x, postfix=""):
 
     # Apply preprocessing factors for all replicas to a given input grid
     def preprocessing_replicas(x, postfix=""):
-        return Lambda(lambda pfs: op.stack(pfs, axis=-1), name=f"prefactors{postfix}")(
+        return Lambda(lambda pfs: op.stack(pfs, axis=1), name=f"prefactors{postfix}")(
             [pf(x) for pf in preprocessing_factor_replicas]
         )
 
@@ -752,7 +752,7 @@ def compute_unnormalized_pdf(x, postfix=""):
             # add batch and flavor dimensions
             photon_integrals = op.batchit(op.batchit(photons.integral))
         else:
-            photon_integrals = np.zeros((1, 1, num_replicas))
+            photon_integrals = np.zeros((1, num_replicas, 1))
 
         PDFs_normalized = sumrule_layer(
             {
@@ -773,7 +773,7 @@ def compute_unnormalized_pdf(x, postfix=""):
     if replica_axis:
         pdf_model = MetaModel(model_input, PDFs, name=f"PDFs", scaler=scaler)
     else:
-        pdf_model = MetaModel(model_input, PDFs[..., 0], name=f"PDFs", scaler=scaler)
+        pdf_model = MetaModel(model_input, PDFs[:, 0], name=f"PDFs", scaler=scaler)
 
     return pdf_model
 
diff --git a/n3fit/src/n3fit/model_trainer.py b/n3fit/src/n3fit/model_trainer.py
index 956cd3a0ea..81af8337ab 100644
--- a/n3fit/src/n3fit/model_trainer.py
+++ b/n3fit/src/n3fit/model_trainer.py
@@ -356,7 +356,7 @@ def _xgrid_generation(self):
         # The PDF model will be called with a concatenation of all inputs
         # now the output needs to be splitted so that each experiment takes its corresponding input
         sp_ar = [[i.shape[1] for i in inputs_unique]]
-        sp_kw = {"axis": 1}
+        sp_kw = {"axis": 2}
         sp_layer = op.as_layer(op.split, op_args=sp_ar, op_kwargs=sp_kw, name="pdf_split")
 
         return InputInfo(input_layer, sp_layer, inputs_idx)
diff --git a/n3fit/src/n3fit/msr.py b/n3fit/src/n3fit/msr.py
index 2c3de11ca6..e781a672d2 100644
--- a/n3fit/src/n3fit/msr.py
+++ b/n3fit/src/n3fit/msr.py
@@ -53,9 +53,9 @@ def generate_msr_model_and_grid(
             - input: the input layer of the integration grid
     """
     # 0. Prepare input layers to MSR model
-    pdf_x = Input(shape=(None, output_dim, replicas), batch_size=1, name="pdf_x")
+    pdf_x = Input(shape=(replicas, None, output_dim), batch_size=1, name="pdf_x")
     pdf_xgrid_integration = Input(
-        shape=(nx, output_dim, replicas), batch_size=1, name="pdf_xgrid_integration"
+        shape=(replicas, nx, output_dim), batch_size=1, name="pdf_xgrid_integration"
     )
 
     # 1. Generate the grid and weights that will be used to integrate
@@ -81,14 +81,14 @@ def generate_msr_model_and_grid(
 
     # 3. Prepare the pdf for integration by dividing by x
     pdf_integrand = Lambda(
-        lambda x_pdf: op.batchit(x_pdf[0], batch_dimension=-1) * x_pdf[1], name="pdf_integrand"
+        lambda x_pdf: op.batchit(x_pdf[0], batch_dimension=1) * x_pdf[1], name="pdf_integrand"
     )([x_divided, pdf_xgrid_integration])
 
     # 4. Integrate the pdf
     pdf_integrated = xIntegrator(weights_array, input_shape=(nx,))(pdf_integrand)
 
     # 5. THe input for the photon integral, will be set to 0 if no photons
-    photon_integral = Input(shape=(1, replicas), batch_size=1, name='photon_integral')
+    photon_integral = Input(shape=(replicas, 1), batch_size=1, name='photon_integral')
 
     # 6. Compute the normalization factor
     normalization_factor = MSR_Normalization(mode, replicas, name="msr_weights")(
diff --git a/n3fit/src/n3fit/stopping.py b/n3fit/src/n3fit/stopping.py
index 5488078a29..8884e7f01b 100644
--- a/n3fit/src/n3fit/stopping.py
+++ b/n3fit/src/n3fit/stopping.py
@@ -349,7 +349,7 @@ def __init__(
         self._positivity = Positivity(threshold_positivity, pos_sets)
 
         # Initialize internal variables for the stopping
-        self._n_replicas = pdf_model.output_shape[-1]
+        self._n_replicas = pdf_model.num_replicas
         self._threshold_chi2 = threshold_chi2
         self._stopping_degrees = np.zeros(self._n_replicas, dtype=int)
         self._counts = np.zeros(self._n_replicas, dtype=int)
diff --git a/n3fit/src/n3fit/tests/test_layers.py b/n3fit/src/n3fit/tests/test_layers.py
index b8cbb844a0..1cae6c6c6a 100644
--- a/n3fit/src/n3fit/tests/test_layers.py
+++ b/n3fit/src/n3fit/tests/test_layers.py
@@ -3,11 +3,12 @@
     This module checks that the layers do what they would do with numpy
 """
 import dataclasses
+
 import numpy as np
-from validphys.pdfbases import fitbasis_to_NN31IC
+
 from n3fit.backends import operations as op
 import n3fit.layers as layers
-
+from validphys.pdfbases import fitbasis_to_NN31IC
 
 FLAVS = 3
 XSIZE = 4
@@ -145,7 +146,7 @@ def test_DIS():
         fks = [i.fktable for i in fktables]
         obs_layer = layers.DIS(fktables, fks, ope, nfl=FLAVS)
         pdf = np.random.rand(XSIZE, FLAVS)
-        kp = op.numpy_to_tensor(np.expand_dims(pdf, 0))
+        kp = op.numpy_to_tensor([[pdf]])  # add batch and replica dimension
         # generate the n3fit results
         result_tensor = obs_layer(kp)
         result = op.evaluate(result_tensor)
@@ -169,8 +170,7 @@ def test_DY():
         fks = [i.fktable for i in fktables]
         obs_layer = layers.DY(fktables, fks, ope, nfl=FLAVS)
         pdf = np.random.rand(XSIZE, FLAVS)
-        # Add batch dimension (0) and replica dimension (-1)
-        kp = op.numpy_to_tensor(np.expand_dims(pdf, [0, -1]))
+        kp = op.numpy_to_tensor([[pdf]])  # add batch and replica dimension
         # generate the n3fit results
         result_tensor = obs_layer(kp)
         result = op.evaluate(result_tensor)
@@ -201,15 +201,15 @@ def test_rotation_flavour():
         {"fl": "g"},
     ]
     # Apply the rotation using numpy tensordot
-    x = np.ones(8)  # Vector in the flavour basis v_i
-    x = np.expand_dims(x, axis=[0, 1])  # Give to the input the shape (1,1,8)
+    pdf = np.ones(8)  # Vector in the flavour basis v_i
+    pdf = np.expand_dims(pdf, axis=[0, 1, 2])  # Add batch, replica, x dimensions
     mat = fitbasis_to_NN31IC(flav_info, "FLAVOUR")  # Rotation matrix R_ij, i=flavour, j=evolution
-    res_np = np.tensordot(x, mat, (2, 0))  # Vector in the evolution basis u_j=R_ij*vi
+    res_np = np.tensordot(pdf, mat, (3, 0))  # Vector in the evolution basis u_j=R_ij*vi
 
     # Apply the rotation through the rotation layer
-    x = op.numpy_to_tensor(x)
+    pdf = op.numpy_to_tensor(pdf)
     rotmat = layers.FlavourToEvolution(flav_info, "FLAVOUR")
-    res_layer = rotmat(x)
+    res_layer = rotmat(pdf)
     assert np.alltrue(res_np == res_layer)
 
 
@@ -226,15 +226,15 @@ def test_rotation_evol():
         {"fl": "g"},
     ]
     # Apply the rotation using numpy tensordot
-    x = np.ones(8)  # Vector in the flavour basis v_i
-    x = np.expand_dims(x, axis=[0, 1])  # Give to the input the shape (1,1,8)
+    pdf = np.ones(8)  # Vector in the flavour basis v_i
+    pdf = np.expand_dims(pdf, axis=[0, 1, 2])  # Add batch, replica, x dimensions
     mat = fitbasis_to_NN31IC(flav_info, "EVOL")  # Rotation matrix R_ij, i=flavour, j=evolution
-    res_np = np.tensordot(x, mat, (2, 0))  # Vector in the evolution basis u_j=R_ij*vi
+    res_np = np.tensordot(pdf, mat, (3, 0))  # Vector in the evolution basis u_j=R_ij*vi
 
     # Apply the rotation through the rotation layer
-    x = op.numpy_to_tensor(x)
+    pdf = op.numpy_to_tensor(pdf)
     rotmat = layers.FlavourToEvolution(flav_info, "EVOL")
-    res_layer = rotmat(x)
+    res_layer = rotmat(pdf)
     assert np.alltrue(res_np == res_layer)
 
 
@@ -260,6 +260,7 @@ def test_mask():
     ret = masker(fi)
     np.testing.assert_allclose(ret, masked_fi * rn_val, rtol=1e-5)
 
+
 def test_addphoton_init():
     """Test AddPhoton class."""
     addphoton = layers.AddPhoton(photons=None)
@@ -268,13 +269,15 @@ def test_addphoton_init():
     np.testing.assert_equal(addphoton._photons_generator, 1234)
     np.testing.assert_equal(addphoton._pdf_ph, None)
 
-class FakePhoton():
+
+class FakePhoton:
     def __call__(self, xgrid):
         return [np.exp(-xgrid)]
 
+
 def test_compute_photon():
     photon = FakePhoton()
     addphoton = layers.AddPhoton(photons=photon)
-    xgrid = np.geomspace(1e-4, 1., 10)
+    xgrid = np.geomspace(1e-4, 1.0, 10)
     addphoton.register_photon(xgrid)
     np.testing.assert_allclose(addphoton._pdf_ph, [np.exp(-xgrid)])
diff --git a/n3fit/src/n3fit/tests/test_msr.py b/n3fit/src/n3fit/tests/test_msr.py
index e51076741c..9f76d3f421 100644
--- a/n3fit/src/n3fit/tests/test_msr.py
+++ b/n3fit/src/n3fit/tests/test_msr.py
@@ -6,9 +6,9 @@
 
 def apply_layer_to_fixed_input(layer):
     np.random.seed(422)
-    pdf_integrated = op.numpy_to_tensor(np.random.normal(size=(1, 14,1)))
+    pdf_integrated = op.numpy_to_tensor(np.random.normal(size=(1, 1, 14)))
 
-    photon_integral = op.numpy_to_tensor([np.random.normal(size=(1,1))])
+    photon_integral = op.numpy_to_tensor([np.random.normal(size=(1, 1))])
 
     return layer(pdf_integrated, photon_integral)
 
@@ -17,25 +17,26 @@ def test_all():
     layer = MSR_Normalization(mode='ALL')
     output = apply_layer_to_fixed_input(layer)
     known_output = op.numpy_to_tensor(
-        np.expand_dims(
+        [
             [
-                1.0,
-                1.0,
-                6.7740397,
-                -3.8735497,
-                15.276561,
-                5.9522753,
-                8.247783,
-                -3.8735497,
-                -3.8735497,
-                1.0,
-                1.0,
-                1.0,
-                1.0,
-                1.0,
-            ],
-            -1,
-        )
+                [
+                    1.0,
+                    1.0,
+                    6.7740397,
+                    -3.8735497,
+                    15.276561,
+                    5.9522753,
+                    8.247783,
+                    -3.8735497,
+                    -3.8735497,
+                    1.0,
+                    1.0,
+                    1.0,
+                    1.0,
+                    1.0,
+                ]
+            ]
+        ]
     )
     np.testing.assert_allclose(output, known_output, rtol=1e-5)
 
@@ -44,25 +45,7 @@ def test_msr():
     layer = MSR_Normalization(mode='MSR')
     output = apply_layer_to_fixed_input(layer)
     known_output = op.numpy_to_tensor(
-        np.expand_dims(
-            [
-                1.0,
-                1.0,
-                6.7740397,
-                1.0,
-                1.0,
-                1.0,
-                1.0,
-                1.0,
-                1.0,
-                1.0,
-                1.0,
-                1.0,
-                1.0,
-                1.0,
-            ],
-            -1,
-        )
+        [[[1.0, 1.0, 6.7740397, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]]]
     )
     np.testing.assert_allclose(output, known_output, rtol=1e-5)
 
@@ -71,24 +54,25 @@ def test_vsr():
     layer = MSR_Normalization(mode='VSR')
     output = apply_layer_to_fixed_input(layer)
     known_output = op.numpy_to_tensor(
-        np.expand_dims(
+        [
             [
-                1.0,
-                1.0,
-                1.0,
-                -3.8735497,
-                15.276561,
-                5.9522753,
-                8.247783,
-                -3.8735497,
-                -3.8735497,
-                1.0,
-                1.0,
-                1.0,
-                1.0,
-                1.0,
-            ],
-            -1,
-        )
+                [
+                    1.0,
+                    1.0,
+                    1.0,
+                    -3.8735497,
+                    15.276561,
+                    5.9522753,
+                    8.247783,
+                    -3.8735497,
+                    -3.8735497,
+                    1.0,
+                    1.0,
+                    1.0,
+                    1.0,
+                    1.0,
+                ]
+            ]
+        ]
     )
     np.testing.assert_allclose(output, known_output, rtol=1e-5)
diff --git a/n3fit/src/n3fit/tests/test_multireplica.py b/n3fit/src/n3fit/tests/test_multireplica.py
index d187560a4e..c3191bef6d 100644
--- a/n3fit/src/n3fit/tests/test_multireplica.py
+++ b/n3fit/src/n3fit/tests/test_multireplica.py
@@ -6,7 +6,7 @@
 def test_replica_split():
     """Check that multi replica pdf and concatenated single output pdfs agree"""
     num_replicas = 3
-    replica_axis = -1
+    replica_axis = 1
     fake_fl = [
         {"fl": i, "largex": [0, 1], "smallx": [1, 2]}
         for i in ["u", "ubar", "d", "dbar", "c", "g", "s", "sbar"]
diff --git a/n3fit/src/n3fit/tests/test_rotations.py b/n3fit/src/n3fit/tests/test_rotations.py
index a639fde249..d10d9153b9 100644
--- a/n3fit/src/n3fit/tests/test_rotations.py
+++ b/n3fit/src/n3fit/tests/test_rotations.py
@@ -8,8 +8,10 @@ def test_fk():
     rotation = FkRotation()
     gridpoints = 2
     np.random.seed(0)
-    pdf = op.numpy_to_tensor(np.random.rand(1, gridpoints, 9))
+    pdf = op.numpy_to_tensor(np.random.rand(1, 1, gridpoints, 9))
     pdf_rotated = rotation(pdf)
+    # extract single replica
+    pdf_rotated = pdf_rotated[:, 0]
     pdf_rotated_known = op.numpy_to_tensor(
         [
             [
diff --git a/n3fit/src/n3fit/tests/test_xops.py b/n3fit/src/n3fit/tests/test_xops.py
index db030ae350..957f573f06 100644
--- a/n3fit/src/n3fit/tests/test_xops.py
+++ b/n3fit/src/n3fit/tests/test_xops.py
@@ -38,10 +38,10 @@ def test_xdivide_indices():
 def test_xintegrator():
     np.random.seed(42)
     weights = np.random.rand(5, 1)
-    pdf = op.numpy_to_tensor(np.random.rand(1, 5, 8))
+    pdf = op.numpy_to_tensor(np.random.rand(1, 1, 5, 8))
     xint = xIntegrator(weights)
     xint_out = xint(pdf)
     xint_out_reference = np.array(
-        [[0.405455, 0.878931, 0.937715, 0.906214, 1.984154, 1.147975, 1.642387, 1.549858]]
+        [[[0.405455, 0.878931, 0.937715, 0.906214, 1.984154, 1.147975, 1.642387, 1.549858]]]
     )
     np.testing.assert_allclose(xint_out.numpy(), xint_out_reference, rtol=1e-05)
diff --git a/validphys2/src/validphys/photon/compute.py b/validphys2/src/validphys/photon/compute.py
index 79a90b8c77..56872540be 100644
--- a/validphys2/src/validphys/photon/compute.py
+++ b/validphys2/src/validphys/photon/compute.py
@@ -187,14 +187,14 @@ def __call__(self, xgrid):
         Returns
         -------
         photon values : nd.array
-            array of photon values with shape (1,xgrid,1,replicas)
+            array of photon values with shape (1, replicas, xgrid, 1)
         """
         return np.stack(
             [
                 self.interpolator[id](xgrid[0, :, 0])[np.newaxis, :, np.newaxis]
                 for id in range(len(self.replicas))
             ],
-            axis=-1,
+            axis=1,
         )
 
     @property
diff --git a/validphys2/src/validphys/tests/photon/test_compute.py b/validphys2/src/validphys/tests/photon/test_compute.py
index 0a204f00fd..0d4c7ea57d 100644
--- a/validphys2/src/validphys/tests/photon/test_compute.py
+++ b/validphys2/src/validphys/tests/photon/test_compute.py
@@ -203,18 +203,8 @@ def test_betas():
     """test betas for different nf"""
     test_theory = API.theoryid(theoryid=THEORY_QED)
     alpha = Alpha(test_theory.get_description(), 1e8)
-    vec_beta0 = [
-        -0.5305164769729844,
-        -0.6719875374991137,
-        -0.7073553026306458,
-        -0.8488263631567751,
-    ]
-    vec_b1 = [
-        0.17507043740108488,
-        0.1605510390839295,
-        0.1538497783221655,
-        0.1458920311675707,
-    ]
+    vec_beta0 = [-0.5305164769729844, -0.6719875374991137, -0.7073553026306458, -0.8488263631567751]
+    vec_b1 = [0.17507043740108488, 0.1605510390839295, 0.1538497783221655, 0.1458920311675707]
     for nf in range(3, 6 + 1):
         np.testing.assert_allclose(alpha.betas_qed[nf][0], vec_beta0[nf - 3], rtol=1e-7)
         np.testing.assert_allclose(alpha.betas_qed[nf][1], vec_b1[nf - 3], rtol=1e-7)
@@ -288,6 +278,6 @@ def test_photon():
         photon_Q0 = pdfs_final[ph_id]
         photon_fiatlux = XGRID * photon_Q0
 
-        photon_validphys = photon(XGRID[np.newaxis, :, np.newaxis])[0, :, 0, 0]
+        photon_validphys = photon(XGRID[np.newaxis, :, np.newaxis])[0, 0, :, 0]
 
         np.testing.assert_allclose(photon_fiatlux, photon_validphys, rtol=1e-7)