From c2e4935a84292397abce3b6b3c811fa668f952ed Mon Sep 17 00:00:00 2001 From: Aron Date: Mon, 19 Feb 2024 13:34:02 +0100 Subject: [PATCH] Fix shape of input and mask in the Mask layer, and adjust test. --- n3fit/src/n3fit/layers/mask.py | 27 ++++++++++++++++++--------- n3fit/src/n3fit/tests/test_layers.py | 9 +++++---- 2 files changed, 23 insertions(+), 13 deletions(-) diff --git a/n3fit/src/n3fit/layers/mask.py b/n3fit/src/n3fit/layers/mask.py index 90c8650499..2b0bff6c00 100644 --- a/n3fit/src/n3fit/layers/mask.py +++ b/n3fit/src/n3fit/layers/mask.py @@ -6,7 +6,7 @@ class Mask(MetaLayer): """ - This layers applies a boolean mask to a rank-1 input tensor. + This layers applies a boolean mask to an input tensor. The mask admit a multiplier for all outputs which will be internally saved as a weight so it can be updated during trainig. @@ -16,7 +16,7 @@ class Mask(MetaLayer): Parameters ---------- - bool_mask: np.array + bool_mask: np.array of shape (n_replicas, n_features) numpy array with the boolean mask to be applied c: float constant multiplier for every output @@ -28,10 +28,7 @@ def __init__(self, bool_mask=None, c=None, **kwargs): self.last_dim = -1 else: self.mask = op.numpy_to_tensor(bool_mask, dtype=bool) - if len(bool_mask.shape) == 1: - self.last_dim = count_nonzero(bool_mask) - else: - self.last_dim = count_nonzero(bool_mask[0, ...]) + self.last_dim = count_nonzero(bool_mask[0, ...]) self.c = c self.masked_output_shape = None super().__init__(**kwargs) @@ -42,12 +39,24 @@ def build(self, input_shape): self.kernel = self.builder_helper("mask", (1,), initializer, trainable=False) # Make sure reshape will succeed: set the last dimension to the unmasked data length and before-last to # the number of replicas - self.masked_output_shape = [-1 if d is None else d for d in input_shape] - self.masked_output_shape[-1] = self.last_dim - self.masked_output_shape[-2] = self.mask.shape[-2] + if self.mask is not None: + self.masked_output_shape = [-1 if d is None else d for d in input_shape] + self.masked_output_shape[-1] = self.last_dim + self.masked_output_shape[-2] = self.mask.shape[-2] super(Mask, self).build(input_shape) def call(self, ret): + """ + Apply the mask to the input tensor, and multiply by the constant if present. + + Parameters + ---------- + ret: Tensor of shape (batch_size, n_replicas, n_features) + + Returns + ------- + Tensor of shape (batch_size, n_replicas, n_features) + """ if self.mask is not None: flat_res = op.boolean_mask(ret, self.mask, axis=1) ret = op.reshape(flat_res, shape=self.masked_output_shape) diff --git a/n3fit/src/n3fit/tests/test_layers.py b/n3fit/src/n3fit/tests/test_layers.py index 1cae6c6c6a..6ff4887d64 100644 --- a/n3fit/src/n3fit/tests/test_layers.py +++ b/n3fit/src/n3fit/tests/test_layers.py @@ -240,8 +240,9 @@ def test_rotation_evol(): def test_mask(): """Test the mask layer""" - SIZE = 100 - fi = np.random.rand(SIZE) + batch_size, replicas, points = 1, 1, 100 + shape = (batch_size, replicas, points) + fi = np.random.rand(*shape) # Check that the multiplier works vals = [0.0, 2.0, np.random.rand()] for val in vals: @@ -249,10 +250,10 @@ def test_mask(): ret = masker(fi) np.testing.assert_allclose(ret, val * fi, rtol=1e-5) # Check that the boolean works - np_mask = np.random.randint(0, 2, size=SIZE, dtype=bool) + np_mask = np.random.randint(0, 2, size=shape[1:], dtype=bool) masker = layers.Mask(bool_mask=np_mask) ret = masker(fi) - masked_fi = fi[np_mask] + masked_fi = fi[np.newaxis, :, np_mask] np.testing.assert_allclose(ret, masked_fi, rtol=1e-5) # Check that the combination works! rn_val = vals[-1]