Remove MultiDropout, not necessary

n3fit/src/n3fit/backends/keras_backend/base_layers.py

@@ -18,13 +18,18 @@
 """
 
 from tensorflow import expand_dims, math, nn
+from tensorflow.keras.layers import (  # pylint: disable=unused-import
+    Dropout,
+    Input,
+    Lambda,
+    concatenate,
+)
 from tensorflow.keras.layers import Dense as KerasDense
-from tensorflow.keras.layers import Input, Lambda, concatenate  # pylint: disable=unused-import
 from tensorflow.keras.layers import LSTM, Concatenate  # pylint: disable=unused-import
 from tensorflow.keras.regularizers import l1_l2
 
 from n3fit.backends import MetaLayer
-from n3fit.backends.keras_backend.multi_dense import MultiDense, MultiDropout
+from n3fit.backends.keras_backend.multi_dense import MultiDense
 
 
 # Custom activation functions
@@ -159,7 +164,7 @@ def apply_dense(xinput):
         LSTM_modified,
         {"kernel_initializer": "glorot_normal", "units": 5, "activation": "sigmoid"},
     ),
-    "dropout": (MultiDropout, {"rate": 0.0}),
+    "dropout": (Dropout, {"rate": 0.0}),
     "concatenate": (Concatenate, {}),
 }
 

n3fit/src/n3fit/backends/keras_backend/multi_dense.py

@@ -177,22 +177,3 @@ def __call__(self, shape, dtype=None, **kwargs):
             per_replica_weights.append(single_initializer(shape, dtype, **kwargs))
 
         return tf.stack(per_replica_weights, axis=0)
-
-
-class MultiDropout(Dropout):
-    """
-    Dropout that broadcasts to the replica axis, to make sure that the dropout rate is constant
-    per replica.
-
-    Input shape: (batch_size, replicas, gridsize, features)
-
-    """
-
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-        self.replica_axis = 1
-
-    def _get_noise_shape(self, inputs):
-        input_shape = list(inputs.shape)
-        noise_shape = input_shape[: self.replica_axis] + [1] + input_shape[self.replica_axis + 1 :]
-        return noise_shape

n3fit/src/n3fit/tests/test_multidense.py

@@ -66,22 +66,3 @@ def test_initializers():
     stacked_weights = stacked_weights.numpy()
 
     np.testing.assert_allclose(multi_dense_weights, stacked_weights)
-
-
-def test_dropout():
-    replicas = 100
-    x_size = 10
-    features = 1
-    input_shape = (1, replicas, x_size, features)
-    test_input = tf.ones(shape=input_shape)
-
-    dropout_layer = MultiDropout(rate=0.5, seed=44)
-
-    test_output = dropout_layer(test_input, training=True)
-
-    # Check that for every replica the same x values are dropped
-    zero_indices_first_replica = np.where(test_output.numpy()[0, 0, :, 0] == 0)
-    nonzero_indices_first_replica = np.where(test_output.numpy()[0, 0, :, 0] != 0)
-
-    assert np.all(test_output.numpy()[:, :, zero_indices_first_replica, :] == 0)
-    assert np.all(test_output.numpy()[:, :, nonzero_indices_first_replica, :] != 0)