Repository: cirKITers/masKIT

Issue:
New PerturbationMode to change parameter values
It would be cool to have a `PertubationMode` that is not changing the number of gates in the circuit but changes the parameter value of gates by adding `pi` to its current value (and thereby inverting the rotation of the gate). 

Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/cirKITers/masKIT/pull/54?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers) Report
> Merging [#54](https://codecov.io/gh/cirKITers/masKIT/pull/54?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers) (7b2e329) into [main](https://codecov.io/gh/cirKITers/masKIT/commit/342378fdfc8affec5e30230628eb85cb5980ec04?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers) (342378f) will **increase** coverage by `0.46%`.
> The diff coverage is `95.45%`.

[![Impacted file tree graph](https://codecov.io/gh/cirKITers/masKIT/pull/54/graphs/tree.svg?width=650&height=150&src=pr&token=8Or2uHjmFE&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers)](https://codecov.io/gh/cirKITers/masKIT/pull/54?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers)

```diff
@@            Coverage Diff             @@
##             main      #54      +/-   ##
==========================================
+ Coverage   81.21%   81.67%   +0.46%     
==========================================
  Files          24       24              
  Lines        1363     1430      +67     
  Branches      155      164       +9     
==========================================
+ Hits         1107     1168      +61     
- Misses        247      250       +3     
- Partials        9       12       +3     
```


| [Impacted Files](https://codecov.io/gh/cirKITers/masKIT/pull/54?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers) | Coverage Δ | |
|---|---|---|
| [main.py](https://codecov.io/gh/cirKITers/masKIT/pull/54/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers#diff-bWFpbi5weQ==) | `0.00% <ø> (ø)` | |
| [tests/configurations.py](https://codecov.io/gh/cirKITers/masKIT/pull/54/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers#diff-dGVzdHMvY29uZmlndXJhdGlvbnMucHk=) | `100.00% <ø> (ø)` | |
| [maskit/\_masks.py](https://codecov.io/gh/cirKITers/masKIT/pull/54/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers#diff-bWFza2l0L19tYXNrcy5weQ==) | `94.69% <91.04%> (-5.31%)` | :arrow_down: |
| [maskit/\_masked\_circuits.py](https://codecov.io/gh/cirKITers/masKIT/pull/54/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers#diff-bWFza2l0L19tYXNrZWRfY2lyY3VpdHMucHk=) | `100.00% <100.00%> (ø)` | |
| [tests/test\_masked\_circuits.py](https://codecov.io/gh/cirKITers/masKIT/pull/54/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers#diff-dGVzdHMvdGVzdF9tYXNrZWRfY2lyY3VpdHMucHk=) | `100.00% <100.00%> (ø)` | |
| [tests/test\_masks.py](https://codecov.io/gh/cirKITers/masKIT/pull/54/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers#diff-dGVzdHMvdGVzdF9tYXNrcy5weQ==) | `100.00% <100.00%> (ø)` | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/cirKITers/masKIT/pull/54?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/cirKITers/masKIT/pull/54?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers). Last update [342378f...7b2e329](https://codecov.io/gh/cirKITers/masKIT/pull/54?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers).

eileen-kuehn: @maxfischer2781, would you please have another look at the issues you pointed out? I tried to cope with all of them :)

Patch:
diff --git a/main.py b/main.py
index 43e0c63..1ae4861 100644
--- a/main.py
+++ b/main.py
@@ -5,7 +5,7 @@ import pennylane as qml
 from pennylane import numpy as np
 
 from maskit._masks import (
-    Mask,
+    DropoutMask,
     PerturbationAxis as Axis,
     PerturbationMode as Mode,
 )
@@ -70,10 +70,10 @@ def init_parameters(
         layers=layers,
         wires=wires,
         default_value=default_value,
-        entangling_mask=Mask(shape=(layers, wires - 1)),
+        entangling_mask=DropoutMask(shape=(layers, wires - 1)),
         dynamic_parameters=dynamic_parameters,
     )
-    mc.mask_for_axis(Axis.LAYERS)[current_layers:] = True
+    mc.mask(Axis.LAYERS)[current_layers:] = True
     return mc
 
 
@@ -173,7 +173,9 @@ def train(
             logging_branch_cost_step["brutto"][step] = result.brutto_steps
             logging_branch_cost_step["netto"][step] = result.netto_steps
         logging_cost_values.append(result.cost)
-        logging_dropout_count_values.append(np.sum(masked_circuit.mask))
+        logging_dropout_count_values.append(
+            np.sum(masked_circuit.full_mask(DropoutMask))
+        )
         if step % train_params["log_interval"] == 0:
             # perform logging
             logging_costs[step] = np.average(logging_cost_values)
@@ -189,7 +191,7 @@ def train(
 
     if __debug__:
         print(masked_circuit.parameters)
-        print(masked_circuit.mask)
+        print(masked_circuit.full_mask(DropoutMask))
 
     return {
         "costs": logging_costs,
@@ -200,10 +202,10 @@ def train(
         "branch_step_costs": logging_branch_cost_step,
         "final_layers": current_layers,
         "params": masked_circuit.parameters.unwrap(),
-        "mask": masked_circuit.mask.unwrap(),
-        "__wire_mask": masked_circuit.mask_for_axis(Axis.WIRES).mask,
-        "__layer_mask": masked_circuit.mask_for_axis(Axis.LAYERS).mask,
-        "__parameter_mask": masked_circuit.mask_for_axis(Axis.PARAMETERS).mask,
+        "dropout_mask": masked_circuit.full_mask(DropoutMask).unwrap(),
+        "__wire_mask": masked_circuit.mask(Axis.WIRES).mask,
+        "__layer_mask": masked_circuit.mask(Axis.LAYERS).mask,
+        "__parameter_mask": masked_circuit.mask(Axis.PARAMETERS).mask,
         "__rotations": rotations,
     }
 
@@ -291,7 +293,7 @@ if __name__ == "__main__":
                     {
                         "perturb": {
                             "amount": 1,
-                            "mode": Mode.ADD,
+                            "mode": Mode.SET,
                             "axis": Axis.PARAMETERS,
                         },
                     },
@@ -301,7 +303,7 @@ if __name__ == "__main__":
                     {
                         "perturb": {
                             "amount": 0.05,
-                            "mode": Mode.REMOVE,
+                            "mode": Mode.RESET,
                             "axis": Axis.PARAMETERS,
                         }
                     },
diff --git a/maskit/_masked_circuits.py b/maskit/_masked_circuits.py
index 17b786b..b48d942 100644
--- a/maskit/_masked_circuits.py
+++ b/maskit/_masked_circuits.py
@@ -1,7 +1,12 @@
 import pennylane.numpy as np
 
 from typing import Dict, Iterable, List, Optional, Tuple, Type, TypeVar, Union
-from maskit._masks import Mask, PerturbationAxis as Axis, PerturbationMode as Mode
+from maskit._masks import (
+    Mask,
+    DropoutMask,
+    PerturbationAxis as Axis,
+    PerturbationMode as Mode,
+)
 
 Self = TypeVar("Self")
 
@@ -23,11 +28,8 @@ class MaskedCircuit(object):
         change size/order
     :param default_value: Default value for gates that are added back in. In case of
         `None` that is also the default, the last known value is assumed
-    :param parameter_mask: Initialization values of paramater mask, defaults to `None`
-    :param layer_mask: Initialization values of layer mask, defaults to `None`
-    :param wire_mask: Initialization values of wire mask, defaults to `None`
-    :param entangling_mask: The mask to apply for entangling gates within the circuit,
-        defaults to None
+    :param masks: List of tuples of `PerturbationAxis` and `Mask` types to specify
+        which masks are supported, defaults to `None`
     """
 
     __slots__ = ("parameters", "default_value", "_dynamic_parameters", "masks")
@@ -51,7 +53,7 @@ class MaskedCircuit(object):
         self.default_value = default_value
         self._dynamic_parameters = dynamic_parameters
 
-        self.masks: Dict[Axis, Mask] = {}
+        self.masks: Dict[Axis, Dict[Type[Mask], Mask]] = {}
         if masks is not None:
             for axis, mask_type in masks:
                 if axis == Axis.PARAMETERS:
@@ -66,24 +68,51 @@ class MaskedCircuit(object):
                     raise NotImplementedError(
                         f"The perturbation {axis} is not supported"
                     )
-                self.masks[axis] = mask_type(
+                self.masks.setdefault(axis, {})[mask_type] = mask_type(
                     shape=shape, parent=self
                 )  # TODO: initial mask is missing
 
     def register_mask(self, axis: Axis, mask: Mask) -> None:
-        assert (
-            axis not in self.masks
-        ), f"Mask for axis {axis} already set: {self.masks[axis]}"
-        self.masks[axis] = mask
+        """
+        Registers a given `mask` for a given ``PerturbationAxis`` and the appropriate
+        type of `Mask`.
+
+        :param axis: The `PerturbationAxis` the mask is applied to
+        :param mask: The mask to register
+        """
+        assert type(mask) not in self.masks.get(
+            axis, {}
+        ), f"{type(mask)} for {axis} already set: {self.masks[axis][type(mask)]}"
+        self.masks.setdefault(axis, {})[type(mask)] = mask
+
+    def mask(self, axis: Axis, mask_type: Type[Mask] = DropoutMask) -> Mask:
+        """
+        Returns the stored mask for a given `axis` and `mask_type`.
+
+        :param axis: The axis the mask is performed on
+        :param mask_type: The type of `Mask` to return
+        """
+        return self.masks[axis][mask_type]
 
-    def mask_for_axis(self, axis: Axis) -> Mask:
-        return self.masks[axis]
+    def _accumulated_mask(self, for_differentiable=True) -> np.ndarray:
+        result = np.zeros(
+            shape=self.parameters.shape,
+            dtype=bool if for_differentiable is True else float,
+        )
+        for mask_type in {
+            key
+            for value in self.masks.values()
+            for key in value
+            if key.relevant_for_differentiation is for_differentiable
+        }:
+            result = result + self.full_mask(mask_type)
+        return result
 
     @property
     def differentiable_parameters(self) -> np.ndarray:
         """Subset of parameters that are not masked and therefore differentiable."""
         if self._dynamic_parameters:
-            return self.parameters[~self.mask]
+            return self.parameters[~self._accumulated_mask()]
         return self.parameters
 
     @differentiable_parameters.setter
@@ -92,33 +121,45 @@ class MaskedCircuit(object):
         Provides a setter for the differentiable parameters. It is ensured that the
         updated values are written into the underlying :py:attr:`~.parameters`.
         """
+        mask = self._accumulated_mask()
         if self._dynamic_parameters:
-            self.parameters[~self.mask] = value
+            self.parameters[~mask] = value
         else:
-            self.parameters[~self.mask] = value[~self.mask]
+            self.parameters[~mask] = value[~mask]
 
-    @property
-    def mask(self) -> np.ndarray:
+    def full_mask(self, mask_type: Type[Mask]) -> np.ndarray:
         """
-        Accumulated mask of layer, wire, and parameter masks.
+        Accumulated mask of layer, wire, and parameter masks for a given type of Mask.
         Note that this mask is readonly.
         """
-        if Axis.PARAMETERS in self.masks:
-            mask = self.masks[Axis.PARAMETERS].mask.copy()
-        else:
-            mask = np.zeros(self.parameters.shape, dtype=bool, requires_grad=False)
-        if Axis.WIRES in self.masks:
-            mask[:, self.masks[Axis.WIRES].mask] = True
-        if Axis.LAYERS in self.masks:
-            mask[self.masks[Axis.LAYERS].mask, :] = True
-        return mask
+        try:
+            result = self.masks[Axis.PARAMETERS][mask_type].mask.copy()
+        except KeyError:
+            result = np.zeros(
+                self.parameters.shape,
+                dtype=mask_type.dtype,
+                requires_grad=False,
+            )
+        shape = result.shape
+        if Axis.WIRES in self.masks and mask_type in self.masks[Axis.WIRES]:
+            wires = self.masks[Axis.WIRES][mask_type].mask
+            # create new axes to match the shape to properly broadcast values
+            result = (
+                result
+                + wires[(np.newaxis, slice(None), *[np.newaxis] * (len(shape) - 2))]
+            )
+        if Axis.LAYERS in self.masks and mask_type in self.masks[Axis.LAYERS]:
+            layers = self.masks[Axis.LAYERS][mask_type].mask
+            # create new axes to match the shape to properly broadcast values
+            result = result + layers[(slice(None), *[np.newaxis] * (len(shape) - 1))]
+        return result
 
     def active(self) -> int:
         """
         Number of active gates in the circuit based on layer, wire, and parameter mask.
         Entangling gates are not included.
         """
-        mask = self.mask
+        mask = self.full_mask(mask_type=DropoutMask)
         return mask.size - np.sum(mask)
 
     def perturb(
@@ -126,6 +167,7 @@ class MaskedCircuit(object):
         axis: Axis = Axis.PARAMETERS,
         amount: Optional[Union[int, float]] = None,
         mode: Mode = Mode.INVERT,
+        mask: Type[Mask] = DropoutMask,
     ):
         """
         Perturbs the MaskedCircuit for a given ``axis`` that is of type
@@ -136,8 +178,9 @@ class MaskedCircuit(object):
         is automatically limited to the actual size of the py:attr:`~.mask`.
 
         :param amount: Number of items to perturb, defaults to None
-        :param axis: Which mask to perturb
+        :param axis: On which axis to perturb
         :param mode: How to perturb, defaults to PerturbationMode.INVERT
+        :param mask: On which mask type to perturb
         :raises NotImplementedError: Raised in case of an unknown mode
         """
         assert mode in list(
@@ -145,21 +188,35 @@ class MaskedCircuit(object):
         ), f"The selected perturbation mode {mode} is not supported."
         if amount == 0:
             return
-        if axis in self.masks:
-            self.masks[axis].perturb(amount=amount, mode=mode)
-        else:  # TODO: change error
-            raise NotImplementedError(f"The perturbation {axis} is not supported")
+        try:
+            self.masks[axis][mask].perturb(amount=amount, mode=mode)
+        except KeyError:
+            raise ValueError(
+                f"The mask {mask} on axis {axis} is not supported"
+            ) from None
+
+    def shrink(
+        self, axis: Axis = Axis.LAYERS, amount: int = 1, mask: Type[Mask] = DropoutMask
+    ):
+        """
+        Shrinks the `mask` by a specified `amount` on a given `axis`.
 
-    def shrink(self, axis: Axis = Axis.LAYERS, amount: int = 1):
-        if axis in self.masks:
-            self.masks[axis].shrink(amount)
-        else:  # TODO: change error
-            raise NotImplementedError(f"The perturbation {axis} is not supported")
+        :param axis: The `PerturbationAxis` the mask is applied on
+        :param amount: The number of elements to shrink
+        :param mask: The type of `Mask` to consider
+        """
+        try:
+            self.masks[axis][mask].shrink(amount)
+        except KeyError:
+            raise ValueError(
+                f"The mask {mask} on axis {axis} is not supported"
+            ) from None
 
     def clear(self):
         """Resets all masks."""
-        for mask in self.masks.values():
-            mask.clear()
+        for masks in self.masks.values():
+            for mask in masks.values():
+                mask.clear()
 
     def apply_mask(self, values: np.ndarray):
         """
@@ -168,7 +225,10 @@ class MaskedCircuit(object):
 
         :param values: Values where the mask should be applied to
         """
-        return values[~self.mask]
+        value_mask = self._accumulated_mask(for_differentiable=False)
+        diff_mask = self._accumulated_mask()
+        result = values + value_mask
+        return result[~diff_mask]
 
     def mask_changed(self, mask: Mask, indices: np.ndarray):
         """
@@ -182,16 +242,17 @@ class MaskedCircuit(object):
             return
         np_indices = tuple(zip(*indices))
         if not np.all(mask.mask[np_indices]):
-            for axis, registered_mask in self.masks.items():
-                if mask == registered_mask:
-                    if axis == Axis.WIRES:
-                        self.parameters[:, np_indices] = self.default_value
-                    elif axis == Axis.LAYERS:
-                        self.parameters[np_indices, :] = self.default_value
-                    elif axis == Axis.PARAMETERS:
-                        self.parameters[np_indices] = self.default_value
-                    else:
-                        raise NotImplementedError(f"The mask {mask} is not supported")
+            for axis, registered_masks in self.masks.items():
+                for registered_mask in registered_masks.values():
+                    if mask == registered_mask:
+                        if axis == Axis.WIRES:
+                            self.parameters[:, np_indices] = self.default_value
+                        elif axis == Axis.LAYERS:
+                            self.parameters[np_indices, :] = self.default_value
+                        elif axis == Axis.PARAMETERS:
+                            self.parameters[np_indices] = self.default_value
+                        else:
+                            raise ValueError(f"The mask {mask} is not registered")
 
     def copy(self: Self) -> Self:
         """Returns a copy of the current MaskedCircuit."""
@@ -200,8 +261,12 @@ class MaskedCircuit(object):
         clone.default_value = self.default_value
         clone._dynamic_parameters = self._dynamic_parameters
         clone.masks = {}
-        for axis, mask in self.masks.items():
-            clone.masks[axis] = mask.copy(clone)
+        for axis, masks in self.masks.items():
+            for mask_type, mask in masks.items():
+                try:
+                    clone.masks[axis][mask_type] = mask.copy(clone)
+                except KeyError:
+                    clone.masks[axis] = {mask_type: mask.copy(clone)}
         return clone
 
     def expanded_parameters(self, changed_parameters: np.ndarray) -> np.ndarray:
@@ -217,11 +282,14 @@ class MaskedCircuit(object):
 
         :param changed_parameters: Current set of differentiable parameters
         """
+        diff_mask = self._accumulated_mask()
+        value_mask = self._accumulated_mask(for_differentiable=False)
         result = self.parameters.astype(object)
         if self._dynamic_parameters:
-            result[~self.mask] = changed_parameters.flatten()
+            result[~diff_mask] = changed_parameters.flatten()
         else:
-            result[~self.mask] = changed_parameters[~self.mask]
+            result[~diff_mask] = changed_parameters[~diff_mask]
+        result = result + value_mask
         return result
 
     @staticmethod
@@ -243,7 +311,7 @@ class MaskedCircuit(object):
         length = 0
         first_layer = True
         result = ["["]
-        for layer, layer_hidden in enumerate(self.mask_for_axis(Axis.LAYERS)):
+        for layer, layer_hidden in enumerate(self.mask(Axis.LAYERS)):
             if first_layer:
                 result.append("[")
                 first_layer = False
@@ -251,9 +319,9 @@ class MaskedCircuit(object):
                 result.append("\n [")
             first_wire = True
             first_value = True
-            for wire, wire_hidden in enumerate(self.mask_for_axis(Axis.WIRES)):
+            for wire, wire_hidden in enumerate(self.mask(Axis.WIRES)):
                 if isinstance(
-                    self.mask_for_axis(Axis.PARAMETERS)[layer][wire].unwrap(),
+                    self.mask(Axis.PARAMETERS)[layer][wire].unwrap(),
                     np.ndarray,
                 ):
                     if first_wire:
@@ -263,7 +331,7 @@ class MaskedCircuit(object):
                         result.append("\n  [")
                     first_value = True
                     for parameter, parameter_hidden in enumerate(
-                        self.mask_for_axis(Axis.PARAMETERS)[layer][wire]
+                        self.mask(Axis.PARAMETERS)[layer][wire]
                     ):
                         if not (layer_hidden or wire_hidden or parameter_hidden):
                             value = format_value(
@@ -282,7 +350,7 @@ class MaskedCircuit(object):
                     if not (
                         layer_hidden
                         or wire_hidden
-                        or self.mask_for_axis(Axis.PARAMETERS)[layer][wire]
+                        or self.mask(Axis.PARAMETERS)[layer][wire]
                     ):
                         value = format_value(self.parameters[layer][wire])
                         length = len(value)
@@ -310,6 +378,7 @@ class MaskedCircuit(object):
         layer_mask: Optional[Mask] = None,
         parameter_mask: Optional[Mask] = None,
     ):
+        """Helper to create a full circuit with `DropoutMasks`."""
         initializable_masks = [axis for axis in Axis if axis != Axis.ENTANGLING]
         for axis, mask in (
             (Axis.WIRES, wire_mask),
@@ -322,148 +391,7 @@ class MaskedCircuit(object):
             parameters=parameters,
             layers=layers,
             wires=wires,
-            masks=[(axis, Mask) for axis in initializable_masks],
-            dynamic_parameters=dynamic_parameters,
-            default_value=default_value,
-        )
-        for axis, mask in (
-            (Axis.WIRES, wire_mask),
-            (Axis.LAYERS, layer_mask),
-            (Axis.PARAMETERS, parameter_mask),
-        ):
-            if mask is not None:
-                circuit.register_mask(axis=axis, mask=mask)
-        if entangling_mask is not None:
-            assert layers == entangling_mask.shape[0]
-            circuit.register_mask(Axis.ENTANGLING, mask=entangling_mask)
-        return circuit
-
-
-class FreezableMaskedCircuit(MaskedCircuit):
-    """
-    A FreezableMaskedCircuit not only supports masking of different components
-    including wires, layers, and parameters but also supports freezing a subset
-    of parameters defined again on the different components wires, layers, and
-    parameters.
-    """
-
-    __slots__ = "freeze_masks"
-
-    def __init__(
-        self,
-        parameters: np.ndarray,
-        layers: int,
-        wires: int,
-        default_value: Optional[float] = None,
-        dynamic_parameters: bool = True,
-        masks: Optional[Iterable[Tuple[Axis, Type[Mask]]]] = None,
-        freeze_masks: Optional[Iterable[Tuple[Axis, Type[Mask]]]] = None,
-    ):
-        super().__init__(
-            parameters,
-            layers,
-            wires,
-            default_value=default_value,
-            dynamic_parameters=dynamic_parameters,
-            masks=masks,
-        )
-        self.freeze_masks: Dict[Axis, Mask] = {}
-        if freeze_masks is not None:
-            for axis, mask_type in freeze_masks:
-                if axis == Axis.PARAMETERS:
-                    shape = parameters.shape
-                elif axis == Axis.WIRES:
-                    shape = (wires,)
-                elif axis == Axis.LAYERS:
-                    shape = (layers,)
-                else:
-                    raise NotImplementedError(
-                        f"The perturbation {axis} is not supported"
-                    )
-                self.freeze_masks[axis] = mask_type(shape=shape)
-
-    @property
-    def mask(self) -> np.ndarray:
-        """
-        Accumulated mask of layers, wires, and parameters for both masking and freezing.
-        Note that this mask is readonly.
-        """
-        base = super().mask
-        if Axis.PARAMETERS in self.freeze_masks:
-            base[self.freeze_masks[Axis.PARAMETERS].mask] = True
-        if Axis.LAYERS in self.freeze_masks:
-            base[self.freeze_masks[Axis.LAYERS].mask, :] = True
-        if Axis.WIRES in self.freeze_masks:
-            base[:, self.freeze_masks[Axis.WIRES].mask] = True
-        return base
-
-    def freeze_mask_for_axis(self, axis: Axis) -> Mask:
-        return self.freeze_masks[axis]
-
-    def freeze(
-        self,
-        axis: Axis = Axis.LAYERS,
-        amount: Optional[Union[int, float]] = None,
-        mode: Mode = Mode.ADD,
-    ):
-        """
-        Freezes the parameter values for a given ``axis`` that is of type
-        :py:class:`~.Axis`. The freezing is applied ``amount``times
-        and depends on the given ``mode`` of type :py:class:`~.PerturbationMode`.
-        If no amount is given, that is ``amount=None``, a random ``amount`` is
-        determined given by the actual size of the py:attr:`~.mask`. The ``amount``
-        is automatically limited to the actual size of the py:attr:`~.mask`.
-
-        :param amount: Number of items to freeze, defaults to None
-        :param axis: Which mask to freeze, defaults to PerturbationAxis.LAYERS
-        :param mode: How to freeze, defaults to PerturbationMode.ADD
-        :raises NotImplementedError: Raised in case of an unknown mode
-        """
-        assert mode in list(
-            Mode
-        ), f"The selected perturbation mode {mode} is not supported."
-        if amount == 0:
-            return
-        if axis in self.freeze_masks:
-            self.freeze_mask_for_axis(axis).perturb(amount=amount, mode=mode)
-        else:
-            raise NotImplementedError(f"The perturbation {axis} is not supported")
-
-    def copy(self: Self) -> Self:
-        """Returns a copy of the current FreezableMaskedCircuit."""
-        clone = super().copy()
-        clone.freeze_masks = {}
-        for axis, mask in self.freeze_masks.items():
-            clone.freeze_masks[axis] = mask.copy()
-        return clone
-
-    @classmethod
-    def full_circuit(
-        cls,
-        parameters: np.ndarray,
-        layers: int,
-        wires: int,
-        dynamic_parameters: bool = True,
-        default_value: Optional[float] = None,
-        entangling_mask: Optional[Mask] = None,
-        wire_mask: Optional[Mask] = None,
-        layer_mask: Optional[Mask] = None,
-        parameter_mask: Optional[Mask] = None,
-    ):
-        initializable_masks = [axis for axis in Axis if axis != Axis.ENTANGLING]
-        for axis, mask in (
-            (Axis.WIRES, wire_mask),
-            (Axis.LAYERS, layer_mask),
-            (Axis.PARAMETERS, parameter_mask),
-        ):
-            if mask is not None:
-                initializable_masks.remove(axis)
-        circuit = cls(
-            parameters=parameters,
-            layers=layers,
-            wires=wires,
-            masks=[(axis, Mask) for axis in initializable_masks],
-            freeze_masks=[(axis, Mask) for axis in Axis if axis != Axis.ENTANGLING],
+            masks=[(axis, DropoutMask) for axis in initializable_masks],
             dynamic_parameters=dynamic_parameters,
             default_value=default_value,
         )
@@ -484,5 +412,5 @@ if __name__ == "__main__":
     parameter = MaskedCircuit(
         np.array(([21, 22, 23], [11, 22, 33], [43, 77, 89])), 3, 3
     )
-    parameter.mask_for_axis(Axis.WIRES)[1] = True
+    parameter.mask(Axis.WIRES)[1] = True
     print(parameter)
diff --git a/maskit/_masks.py b/maskit/_masks.py
index 0cf20f1..e77b653 100644
--- a/maskit/_masks.py
+++ b/maskit/_masks.py
@@ -1,11 +1,14 @@
 import random as rand
 import pennylane.numpy as np
+from abc import abstractmethod
 from enum import Enum
-from typing import TYPE_CHECKING, Any, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Generic, Optional, Tuple, TypeVar, Union, Type
 
 if TYPE_CHECKING:
     from maskit._masked_circuits import MaskedCircuit
 
+T = TypeVar("T")
+
 
 class PerturbationAxis(Enum):
     #: Perturbation affects whole wires
@@ -19,15 +22,15 @@ class PerturbationAxis(Enum):
 
 
 class PerturbationMode(Enum):
-    #: Adding new holes to the mask
-    ADD = 0
-    #: Removing holes from the mask
-    REMOVE = 1
+    #: Set a new value on a mask
+    SET = 0
+    #: Remove a value from the mask
+    RESET = 1
     #: Invert current state of the mask
     INVERT = 2
 
 
-class Mask(object):
+class Mask(Generic[T]):
     """
     A Mask encapsulates a :py:attr:`~.mask` storing boolean value if a specific value
     is masked or not. In case a specific position is `True`, the according value is
@@ -38,6 +41,9 @@ class Mask(object):
     :param mask: Preset of values that is taken by mask
     """
 
+    relevant_for_differentiation: bool = False
+    #: The underlying datatype of the encapsulated mask
+    dtype: Type[T]
     __slots__ = ("mask", "_parent")
 
     def __init__(
@@ -46,11 +52,12 @@ class Mask(object):
         parent: Optional["MaskedCircuit"] = None,
         mask: Optional[np.ndarray] = None,
     ):
-        super().__init__()
-        self.mask = np.zeros(shape, dtype=bool, requires_grad=False)
+        self.mask = np.zeros(shape, dtype=self.dtype, requires_grad=False)
         if mask is not None:
-            assert mask.dtype == bool, "Mask must be of type bool"
-            assert mask.shape == shape, "Shape of mask must be equal to shape"
+            assert mask.dtype == self.dtype, f"Mask must be of type {self.dtype}"
+            assert (
+                mask.shape == shape
+            ), f"Shape of mask ({mask.shape}) must be equal to {shape}"
             self.mask[:] = mask
         self._parent = parent
 
@@ -68,7 +75,7 @@ class Mask(object):
         """Returns the size of the encapsulated :py:attr:`~.mask`"""
         return self.mask.size
 
-    def __setitem__(self, key, value: bool):
+    def __setitem__(self, key, value: T):
         """
         Convenience function to set the value of a specific position of the
         encapsulated :py:attr:`~.mask`.
@@ -88,11 +95,13 @@ class Mask(object):
                 mask[2, 2] = True
         """
         if isinstance(key, int) or isinstance(key, slice) or isinstance(key, tuple):
-            before = self.mask.copy()
-            self.mask[key] = value
-            delta_indices = np.argwhere(before != self.mask)
             if self._parent is not None:
+                before = self.mask.copy()
+                self.mask[key] = value
+                delta_indices = np.argwhere(before != self.mask)
                 self._parent.mask_changed(self, delta_indices)
+            else:
+                self.mask[key] = value
         else:
             raise NotImplementedError(f"key {key}")
 
@@ -105,19 +114,21 @@ class Mask(object):
             return self.mask[key]
         raise NotImplementedError(f"key {key}")
 
-    def apply_mask(self, values: np.ndarray):
+    @abstractmethod
+    def apply_mask(self, values: np.ndarray) -> np.ndarray:
         """
         Applies the encapsulated py:attr:`~.mask` to the given ``values``.
         Note that the values should have the same shape as the py:attr:`~.mask`.
 
         :param values: Values where the mask should be applied to
         """
-        return values[~self.mask]
+        return NotImplemented
 
     def clear(self) -> None:
         """Resets the mask to not mask anything."""
-        self.mask = np.zeros_like(self.mask, dtype=bool, requires_grad=False)
+        self.mask = np.zeros_like(self.mask, dtype=self.mask.dtype, requires_grad=False)
 
+    @abstractmethod
     def perturb(
         self,
         amount: Optional[Union[int, float]] = None,
@@ -137,6 +148,38 @@ class Mask(object):
         :param mode: How to perturb, defaults to PerturbationMode.INVERT
         :raises NotImplementedError: Raised in case of an unknown mode
         """
+        return NotImplemented
+
+    @abstractmethod
+    def shrink(self, amount: int = 1):
+        """
+        Shrinks a Mask for ``amount`` given positions. Shrinking can be interpreted
+        as setting the ``amount`` leftmost values of the Mask that are different
+        from its default to its default value, e.g. to `False` in case of a boolean
+        `dtype` or `0` in case of floats.
+
+        :param amount: Number of items to reset to `0`, defaults to `1`
+        """
+        return NotImplemented
+
+    def copy(self, parent: Optional["MaskedCircuit"] = None) -> "Mask":
+        """Returns a copy of the current Mask."""
+        clone = object.__new__(type(self))
+        clone.mask = self.mask.copy()
+        clone._parent = parent
+        return clone
+
+    def _count_for_amount(self, amount: Optional[Union[int, float]]) -> int:
+        """
+        Returns a valid number of elements with respect to the current size of the
+        encapsulated :py:attr:`~.mask`. The number of elements is dependent on the
+        value of ``amount``: 1) either a random number is returned in case of `None`;
+        2) a specific percentage in case of an amount in the range of `[0, 1[`;
+        or 3) the absolute value is given.
+
+        :param amount: Amount of elements to consider, random if `None`, percentage
+            if in the range `[0, 1[`, or absolute value otherwise
+        """
         assert (
             amount is None or amount >= 0
         ), "Negative values are not supported, please use PerturbationMode.REMOVE"
@@ -144,14 +187,39 @@ class Mask(object):
             if amount < 1:
                 amount *= self.mask.size
             amount = round(amount)
-        count = abs(amount) if amount is not None else rand.randrange(0, self.mask.size)
+        return (
+            min(abs(amount), self.mask.size)
+            if amount is not None
+            else rand.randrange(0, self.mask.size)
+        )
+
+
+class DropoutMask(Mask[bool]):
+    """
+    A DropoutMask marks the positions that should be dropped from a given array
+    and therefore an associated gate in a quantum circuit.
+    A position that will not be dropped is marked as `False`, and `True` otherwise.
+    """
+
+    dtype = bool
+    relevant_for_differentiation = True
+
+    def apply_mask(self, values: np.ndarray):
+        return values[~self.mask]
+
+    def perturb(
+        self,
+        amount: Optional[Union[int, float]] = None,
+        mode: PerturbationMode = PerturbationMode.INVERT,
+    ):
+        count = self._count_for_amount(amount=amount)
         if count == 0:
             return
-        if mode == PerturbationMode.ADD:
+        if mode == PerturbationMode.SET:
             indices = np.argwhere(~self.mask)
         elif mode == PerturbationMode.INVERT:
             indices = np.array([list(index) for index in np.ndindex(*self.mask.shape)])
-        elif mode == PerturbationMode.REMOVE:
+        elif mode == PerturbationMode.RESET:
             indices = np.argwhere(self.mask)
         else:
             raise NotImplementedError(f"The perturbation mode {mode} is not supported")
@@ -174,9 +242,62 @@ class Mask(object):
         if index.size > 0:
             self[tuple(zip(*index))] = False
 
-    def copy(self, parent: Optional["MaskedCircuit"] = None) -> "Mask":
-        """Returns a copy of the current Mask."""
-        clone = object.__new__(type(self))
-        clone.mask = self.mask.copy()
-        clone._parent = parent
-        return clone
+
+class FreezeMask(DropoutMask):
+    """
+    A FreezeMask provides the same functionality as a :py:class:`~.DropoutMask`.
+    However, the meaning is a bit different. Marked positions are interpreted as
+    being frozen and underlying values therefore are not intended to be changed.
+    """
+
+    pass
+
+
+class ValueMask(Mask[float]):
+    """
+    A ValueMask shifts given values the mask is applied to with the stored values
+    inside the mask.
+    """
+
+    dtype = float
+    relevant_for_differentiation = False
+
+    def apply_mask(self, values: np.ndarray):
+        return values + self.mask
+
+    def perturb(
+        self,
+        amount: Optional[Union[int, float]] = None,
+        mode: PerturbationMode = PerturbationMode.INVERT,
+        value: Optional[float] = np.pi,
+    ):
+        count = self._count_for_amount(amount=amount)
+        if count == 0:
+            return
+        if mode == PerturbationMode.SET:  # INVERT might be added here
+            indices = np.array([list(index) for index in np.ndindex(*self.mask.shape)])
+        elif mode == PerturbationMode.RESET:
+            indices = np.argwhere(self.mask != 0)
+        else:
+            raise NotImplementedError(f"The perturbation mode {mode} is not supported")
+        if indices.size == 0:
+            return
+        indices = tuple(
+            zip(
+                *indices[
+                    np.random.choice(
+                        len(indices), min(count, len(indices)), replace=False
+                    )
+                ]
+            )
+        )
+        if mode == PerturbationMode.SET and value is not None:
+            self[indices] = value
+        elif mode == PerturbationMode.RESET:
+            self[indices] = 0
+
+    def shrink(self, amount: int = 1):
+        index = np.argwhere(self.mask != 0)
+        index = index[:amount]
+        if index.size > 0:
+            self[tuple(zip(*index))] = 0
diff --git a/maskit/circuits.py b/maskit/circuits.py
index 0395f2a..3df04ff 100644
--- a/maskit/circuits.py
+++ b/maskit/circuits.py
@@ -1,4 +1,4 @@
-from maskit._masks import PerturbationAxis as Axis
+from maskit._masks import PerturbationAxis as Axis, DropoutMask
 from maskit._masked_circuits import MaskedCircuit
 import pennylane as qml
 from pennylane import numpy as np
@@ -6,17 +6,15 @@ from pennylane import numpy as np
 
 def basic_variational_circuit(params, rotations, masked_circuit: MaskedCircuit):
     full_parameters = masked_circuit.expanded_parameters(params)
-    wires = len(masked_circuit.mask_for_axis(Axis.WIRES))
-    mask = masked_circuit.mask
-    for wire, _is_masked in enumerate(masked_circuit.mask_for_axis(Axis.WIRES)):
+    wires = len(masked_circuit.mask(Axis.WIRES))
+    dropout_mask = masked_circuit.full_mask(DropoutMask)
+    for wire, _is_masked in enumerate(masked_circuit.mask(Axis.WIRES)):
         qml.RY(np.pi / 4, wires=wire)
     r = -1
-    for layer, _is_layer_masked in enumerate(masked_circuit.mask_for_axis(Axis.LAYERS)):
-        for wire, _is_wire_masked in enumerate(
-            masked_circuit.mask_for_axis(Axis.WIRES)
-        ):
+    for layer, _is_layer_masked in enumerate(masked_circuit.mask(Axis.LAYERS)):
+        for wire, _is_wire_masked in enumerate(masked_circuit.mask(Axis.WIRES)):
             r += 1
-            if mask[layer][wire]:
+            if dropout_mask[layer][wire]:
                 continue
             if rotations[r] == 0:
                 rotation = qml.RX
@@ -28,15 +26,15 @@ def basic_variational_circuit(params, rotations, masked_circuit: MaskedCircuit):
 
         for wire in range(0, wires - 1, 2):
             if (
-                masked_circuit.mask_for_axis(Axis.ENTANGLING) is not None
-                and masked_circuit.mask_for_axis(Axis.ENTANGLING)[layer, wire]
+                masked_circuit.mask(Axis.ENTANGLING) is not None
+                and masked_circuit.mask(Axis.ENTANGLING)[layer, wire]
             ):
                 continue
             qml.CZ(wires=[wire, wire + 1])
         for wire in range(1, wires - 1, 2):
             if (
-                masked_circuit.mask_for_axis(Axis.ENTANGLING) is not None
-                and masked_circuit.mask_for_axis(Axis.ENTANGLING)[layer, wire]
+                masked_circuit.mask(Axis.ENTANGLING) is not None
+                and masked_circuit.mask(Axis.ENTANGLING)[layer, wire]
             ):
                 continue
             qml.CZ(wires=[wire, wire + 1])
@@ -49,7 +47,9 @@ def variational_circuit(params, rotations, masked_circuit):
         rotations=rotations,
         masked_circuit=masked_circuit,
     )
-    return qml.probs(wires=range(len(masked_circuit.mask_for_axis(Axis.WIRES))))
+    return qml.probs(
+        wires=range(len(masked_circuit.mask(axis=Axis.WIRES, mask_type=DropoutMask)))
+    )
 
 
 def iris_circuit(params, data, rotations, masked_circuit):


Test patch:
diff --git a/tests/configurations.py b/tests/configurations.py
index b83a7fb..d162434 100644
--- a/tests/configurations.py
+++ b/tests/configurations.py
@@ -7,7 +7,7 @@ CLASSICAL = {
         {
             "perturb": {
                 "amount": 0.1,
-                "mode": PerturbationMode.ADD,
+                "mode": PerturbationMode.SET,
                 "axis": PerturbationAxis.PARAMETERS,
             }
         },
@@ -23,7 +23,7 @@ RANDOM = {
         {
             "perturb": {
                 "amount": 1,
-                "mode": PerturbationMode.REMOVE,
+                "mode": PerturbationMode.RESET,
                 "axis": PerturbationAxis.PARAMETERS,
             }
         },
@@ -47,7 +47,7 @@ QHACK = {
         {
             "perturb": {
                 "amount": 1,
-                "mode": PerturbationMode.ADD,
+                "mode": PerturbationMode.SET,
                 "axis": PerturbationAxis.PARAMETERS,
             },
         },
@@ -57,7 +57,7 @@ QHACK = {
         {
             "perturb": {
                 "amount": 0.05,
-                "mode": PerturbationMode.REMOVE,
+                "mode": PerturbationMode.RESET,
                 "axis": PerturbationAxis.PARAMETERS,
             }
         },
diff --git a/tests/test_ensembles.py b/tests/test_ensembles.py
index e63f35c..da8cf50 100644
--- a/tests/test_ensembles.py
+++ b/tests/test_ensembles.py
@@ -35,9 +35,7 @@ class TestEnsemble:
     def test_ensemble_step(self, dropout):
         mp = create_circuit(3, layer_size=2)
         optimizer = ExtendedGradientDescentOptimizer()
-        circuit = qml.QNode(
-            variational_circuit, device(mp.mask_for_axis(Axis.WIRES).size)
-        )
+        circuit = qml.QNode(variational_circuit, device(mp.mask(Axis.WIRES).size))
         ensemble = Ensemble(dropout=dropout)
 
         def cost_fn(params, masked_circuit=None):
@@ -74,9 +72,7 @@ class TestIntervalEnsemble:
         interval = 3
         mp = create_circuit(3, layer_size=2)
         optimizer = ExtendedGradientDescentOptimizer()
-        circuit = qml.QNode(
-            variational_circuit, device(mp.mask_for_axis(Axis.WIRES).size)
-        )
+        circuit = qml.QNode(variational_circuit, device(mp.mask(Axis.WIRES).size))
         simple_ensemble = Ensemble(dropout=None)
         interval_ensemble = IntervalEnsemble(dropout=dropout, interval=interval)
 
@@ -124,9 +120,7 @@ class TestAdaptiveEnsemble:
         pnp.random.seed(1234)
         mp = create_circuit(3, layer_size=2)
         optimizer = ExtendedGradientDescentOptimizer()
-        circuit = qml.QNode(
-            variational_circuit, device(mp.mask_for_axis(Axis.WIRES).size)
-        )
+        circuit = qml.QNode(variational_circuit, device(mp.mask(Axis.WIRES).size))
         simple_ensemble = Ensemble(dropout=None)
         adaptive_ensemble = AdaptiveEnsemble(dropout=dropout, size=3, epsilon=0.01)
 
@@ -166,9 +160,7 @@ class TestEnsembleUseCases:
         pnp.random.seed(1234)
         mp = create_circuit(3, layer_size=2)
         ensemble = Ensemble(dropout=CLASSICAL)
-        circuit = qml.QNode(
-            variational_circuit, device(mp.mask_for_axis(Axis.WIRES).size)
-        )
+        circuit = qml.QNode(variational_circuit, device(mp.mask(Axis.WIRES).size))
         optimizer = ExtendedGradientDescentOptimizer()
 
         def cost_fn(params, masked_circuit=None):
@@ -189,11 +181,9 @@ class TestEnsembleUseCases:
         random.seed(1234)
         pnp.random.seed(1234)
         mp = create_circuit(3, layer_size=2)
-        mp.mask_for_axis(Axis.LAYERS)[1:] = True
+        mp.mask(Axis.LAYERS)[1:] = True
         ensemble = Ensemble(dropout=GROWING)
-        circuit = qml.QNode(
-            variational_circuit, device(mp.mask_for_axis(Axis.WIRES).size)
-        )
+        circuit = qml.QNode(variational_circuit, device(mp.mask(Axis.WIRES).size))
         optimizer = ExtendedGradientDescentOptimizer()
 
         def cost_fn(params, masked_circuit=None):
@@ -204,22 +194,20 @@ class TestEnsembleUseCases:
             )
 
         current_cost = 1.0
-        assert pnp.sum(mp.mask_for_axis(Axis.LAYERS)) == 2
-        for _ in range(len(mp.mask_for_axis(Axis.LAYERS)) - 1):
+        assert pnp.sum(mp.mask(Axis.LAYERS)) == 2
+        for _ in range(len(mp.mask(Axis.LAYERS)) - 1):
             result = ensemble.step(mp, optimizer, cost_fn)
             mp = result.branch
             current_cost = result.cost
         assert current_cost == pytest.approx(0.86318044)
-        assert pnp.sum(mp.mask_for_axis(Axis.LAYERS)) == 0
+        assert pnp.sum(mp.mask(Axis.LAYERS)) == 0
 
     def test_random(self):
         random.seed(1234)
         pnp.random.seed(1234)
         mp = create_circuit(3, layer_size=2)
         ensemble = Ensemble(dropout=RANDOM)
-        circuit = qml.QNode(
-            variational_circuit, device(mp.mask_for_axis(Axis.WIRES).size)
-        )
+        circuit = qml.QNode(variational_circuit, device(mp.mask(Axis.WIRES).size))
         optimizer = ExtendedGradientDescentOptimizer()
 
         def cost_fn(params, masked_circuit=None):
@@ -241,9 +229,7 @@ class TestEnsembleUseCases:
         pnp.random.seed(1234)
         mp = create_circuit(3, layer_size=2)
         ensemble = Ensemble(dropout=QHACK)
-        circuit = qml.QNode(
-            variational_circuit, device(mp.mask_for_axis(Axis.WIRES).size)
-        )
+        circuit = qml.QNode(variational_circuit, device(mp.mask(Axis.WIRES).size))
         optimizer = ExtendedGradientDescentOptimizer()
 
         def cost_fn(params, masked_circuit=None):
diff --git a/tests/test_masked_circuits.py b/tests/test_masked_circuits.py
index a33155d..302a8ca 100644
--- a/tests/test_masked_circuits.py
+++ b/tests/test_masked_circuits.py
@@ -5,8 +5,13 @@ import pennylane as qml
 import pennylane.numpy as pnp
 
 from maskit.circuits import variational_circuit as original_variational_circuit
-from maskit._masks import PerturbationAxis as Axis, PerturbationMode as Mode, Mask
-from maskit._masked_circuits import FreezableMaskedCircuit, MaskedCircuit
+from maskit._masks import (
+    PerturbationAxis as Axis,
+    PerturbationMode as Mode,
+    DropoutMask,
+    FreezeMask,
+)
+from maskit._masked_circuits import MaskedCircuit
 
 from tests.utils import cost, create_freezable_circuit, device, variational_circuit
 
@@ -15,9 +20,9 @@ class TestMaskedCircuits:
     def test_init(self):
         mp = self._create_circuit(3)
         assert mp
-        assert len(mp.mask_for_axis(Axis.WIRES)) == 3
-        assert len(mp.mask_for_axis(Axis.LAYERS)) == 3
-        assert mp.mask_for_axis(Axis.PARAMETERS).shape == (3, 3)
+        assert len(mp.mask(Axis.WIRES)) == 3
+        assert len(mp.mask(Axis.LAYERS)) == 3
+        assert mp.mask(Axis.PARAMETERS).shape == (3, 3)
 
         size = 3
         with pytest.raises(AssertionError):
@@ -49,24 +54,22 @@ class TestMaskedCircuits:
                 parameters=pnp.random.uniform(low=0, high=1, size=(size, size)),
                 layers=size,
                 wires=size,
-                entangling_mask=Mask(shape=(size + 1, size)),
+                entangling_mask=DropoutMask(shape=(size + 1, size)),
             )
         with pytest.raises(NotImplementedError):
             MaskedCircuit(
                 parameters=pnp.random.uniform(low=0, high=1, size=(size, size)),
                 layers=size,
                 wires=size,
-                masks=[(Axis.ENTANGLING, Mask)],
+                masks=[(Axis.ENTANGLING, DropoutMask)],
             )
         mc = MaskedCircuit.full_circuit(
             parameters=pnp.random.uniform(low=0, high=1, size=(size, size)),
             layers=size,
             wires=size,
-            wire_mask=pnp.ones((size,), dtype=bool),
-        )
-        assert pnp.array_equal(
-            mc.mask_for_axis(Axis.WIRES), pnp.ones((size,), dtype=bool)
+            wire_mask=DropoutMask(shape=(size,), mask=pnp.ones((size,), dtype=bool)),
         )
+        assert pnp.array_equal(mc.mask(Axis.WIRES), pnp.ones((size,), dtype=bool))
 
     def test_mask(self):
         """
@@ -81,7 +84,7 @@ class TestMaskedCircuits:
             layers=size,
             wires=size,
         )
-        assert mc.mask.size == size * size
+        assert mc.full_mask(DropoutMask).size == size * size
 
         # test circuit with no masks
         mc = MaskedCircuit(
@@ -89,16 +92,16 @@ class TestMaskedCircuits:
             layers=size,
             wires=size,
         )
-        assert mc.mask.size == size * size
+        assert mc.full_mask(DropoutMask).size == size * size
 
         # test circuit containing only layer mask
         mc = MaskedCircuit(
             parameters=pnp.random.uniform(low=0, high=1, size=(size, size)),
             layers=size,
             wires=size,
-            masks=[(Axis.LAYERS, Mask)],
+            masks=[(Axis.LAYERS, DropoutMask)],
         )
-        assert mc.mask.size == size * size
+        assert mc.full_mask(DropoutMask).size == size * size
 
     def test_wrong_mode(self):
         mp = self._create_circuit_with_entangling_gates(3)
@@ -107,7 +110,7 @@ class TestMaskedCircuits:
 
     def test_wrong_axis(self):
         mp = self._create_circuit_with_entangling_gates(3)
-        with pytest.raises(NotImplementedError):
+        with pytest.raises(ValueError):
             mp.perturb(axis=10)
 
     @pytest.mark.parametrize("axis", list(Axis))
@@ -116,96 +119,94 @@ class TestMaskedCircuits:
         mp = self._create_circuit_with_entangling_gates(size)
         mp.perturb(axis=axis, amount=size)
         assert (
-            pnp.sum(mp.mask_for_axis(Axis.LAYERS))
-            + pnp.sum(mp.mask_for_axis(Axis.WIRES))
-            + pnp.sum(mp.mask_for_axis(Axis.PARAMETERS))
-            + pnp.sum(mp.mask_for_axis(Axis.ENTANGLING))
+            pnp.sum(mp.mask(Axis.LAYERS))
+            + pnp.sum(mp.mask(Axis.WIRES))
+            + pnp.sum(mp.mask(Axis.PARAMETERS))
+            + pnp.sum(mp.mask(Axis.ENTANGLING))
             == size
         )
         mp.clear()
         assert (
-            pnp.sum(mp.mask_for_axis(Axis.LAYERS))
-            + pnp.sum(mp.mask_for_axis(Axis.WIRES))
-            + pnp.sum(mp.mask_for_axis(Axis.PARAMETERS))
-            + pnp.sum(mp.mask_for_axis(Axis.ENTANGLING))
+            pnp.sum(mp.mask(Axis.LAYERS))
+            + pnp.sum(mp.mask(Axis.WIRES))
+            + pnp.sum(mp.mask(Axis.PARAMETERS))
+            + pnp.sum(mp.mask(Axis.ENTANGLING))
             == 0
         )
 
     def test_perturb_entangling(self):
         size = 3
         mp = self._create_circuit(size)
-        with pytest.raises(NotImplementedError):
+        with pytest.raises(ValueError):
             mp.perturb(axis=Axis.ENTANGLING, amount=1)
         # ensure nothing has happened as entangling mask is None
         assert (
-            pnp.sum(mp.mask_for_axis(Axis.LAYERS))
-            + pnp.sum(mp.mask_for_axis(Axis.WIRES))
-            + pnp.sum(mp.mask_for_axis(Axis.PARAMETERS))
+            pnp.sum(mp.mask(Axis.LAYERS))
+            + pnp.sum(mp.mask(Axis.WIRES))
+            + pnp.sum(mp.mask(Axis.PARAMETERS))
             == 0
         )
 
         mp = self._create_circuit_with_entangling_gates(size)
         mp.perturb(axis=Axis.ENTANGLING, amount=1)
-        assert pnp.sum(mp.mask_for_axis(Axis.ENTANGLING)) == 1
+        assert pnp.sum(mp.mask(Axis.ENTANGLING)) == 1
 
     @pytest.mark.parametrize("axis", list(Axis))
     def test_zero_amount(self, axis):
         mp = self._create_circuit_with_entangling_gates(3)
         pre_sum = (
-            pnp.sum(mp.mask_for_axis(Axis.WIRES))
-            + pnp.sum(mp.mask_for_axis(Axis.LAYERS))
-            + pnp.sum(mp.mask_for_axis(Axis.PARAMETERS))
-            + pnp.sum(mp.mask_for_axis(Axis.ENTANGLING))
+            pnp.sum(mp.mask(Axis.WIRES))
+            + pnp.sum(mp.mask(Axis.LAYERS))
+            + pnp.sum(mp.mask(Axis.PARAMETERS))
+            + pnp.sum(mp.mask(Axis.ENTANGLING))
         )
         mp.perturb(axis=axis, amount=0)
-        assert pre_sum == pnp.sum(mp.mask_for_axis(Axis.WIRES)) + pnp.sum(
-            mp.mask_for_axis(Axis.LAYERS)
-        ) + pnp.sum(mp.mask_for_axis(Axis.PARAMETERS)) + pnp.sum(
-            mp.mask_for_axis(Axis.ENTANGLING)
-        )
+        assert pre_sum == pnp.sum(mp.mask(Axis.WIRES)) + pnp.sum(
+            mp.mask(Axis.LAYERS)
+        ) + pnp.sum(mp.mask(Axis.PARAMETERS)) + pnp.sum(mp.mask(Axis.ENTANGLING))
 
     def test_apply_mask(self):
         size = 3
         mp = self._create_circuit(size)
-        with pytest.raises(IndexError):
+        with pytest.raises(ValueError):
             mp.apply_mask(pnp.ones((size, size - 1)))
-        mp.mask_for_axis(Axis.WIRES)[: size - 1] = True
+        mp.mask(Axis.WIRES)[: size - 1] = True
         result = mp.apply_mask(pnp.ones((size, size), dtype=bool))
-        assert pnp.sum(~mp.mask) == size
+        assert pnp.sum(~mp.full_mask(DropoutMask)) == size
         assert pnp.sum(result) == size
-        mp.mask_for_axis(Axis.LAYERS)[: size - 1] = True
+        mp.mask(Axis.LAYERS)[: size - 1] = True
         result = mp.apply_mask(pnp.ones((size, size), dtype=bool))
-        assert pnp.sum(~mp.mask) == 1
+        assert pnp.sum(~mp.full_mask(DropoutMask)) == 1
         assert pnp.sum(result) == 1
-        mp.mask_for_axis(Axis.PARAMETERS)[(size - 1, size - 1)] = True
+        mp.mask(Axis.PARAMETERS)[(size - 1, size - 1)] = True
         result = mp.apply_mask(pnp.ones((size, size), dtype=bool))
-        assert pnp.sum(~mp.mask) == 0
+        assert pnp.sum(~mp.full_mask(DropoutMask)) == 0
         assert pnp.sum(result) == 0
 
     def test_copy(self):
         size = 3
         mp = self._create_circuit(size)
-        mp.mask_for_axis(Axis.LAYERS)[0] = True
+        mp.mask(Axis.LAYERS)[0] = True
         new_mp = mp.copy()
 
-        mp.mask_for_axis(Axis.WIRES)[0] = True
-        mp.mask_for_axis(Axis.PARAMETERS)[0, 0] = True
+        mp.mask(Axis.WIRES)[0] = True
+        mp.mask(Axis.PARAMETERS)[0, 0] = True
 
-        assert pnp.sum(mp.mask) > pnp.sum(new_mp.mask)
-        assert pnp.sum(new_mp.mask_for_axis(Axis.WIRES)) == 0
-        assert pnp.sum(new_mp.mask_for_axis(Axis.LAYERS)) == pnp.sum(
-            mp.mask_for_axis(Axis.LAYERS)
+        assert pnp.sum(mp.full_mask(DropoutMask)) > pnp.sum(
+            new_mp.full_mask(DropoutMask)
         )
-        assert pnp.sum(new_mp.mask_for_axis(Axis.PARAMETERS)) == 0
+        assert pnp.sum(new_mp.mask(Axis.WIRES)) == 0
+        assert pnp.sum(new_mp.mask(Axis.LAYERS)) == pnp.sum(mp.mask(Axis.LAYERS))
+        assert pnp.sum(new_mp.mask(Axis.PARAMETERS)) == 0
         assert Axis.ENTANGLING not in new_mp.masks
 
         # also test copying of existing entanglement mask
         mp = self._create_circuit_with_entangling_gates(size)
-        assert mp.mask_for_axis(Axis.ENTANGLING) is not None
+        assert mp.mask(Axis.ENTANGLING) is not None
         new_mp = mp.copy()
-        mp.mask_for_axis(Axis.ENTANGLING)[0, 0] = True
-        assert pnp.sum(mp.mask_for_axis(Axis.ENTANGLING)) == 1
-        assert pnp.sum(new_mp.mask_for_axis(Axis.ENTANGLING)) == 0
+        mp.mask(Axis.ENTANGLING)[0, 0] = True
+        assert pnp.sum(mp.mask(Axis.ENTANGLING)) == 1
+        assert pnp.sum(new_mp.mask(Axis.ENTANGLING)) == 0
 
     def test_parameters(self):
         size = 3
@@ -221,44 +222,47 @@ class TestMaskedCircuits:
     def test_shrink_layer(self):
         size = 3
         mp = self._create_circuit(size)
-        mp.mask_for_axis(Axis.LAYERS)[:] = True
+        mp.mask(Axis.LAYERS)[:] = True
         mp.shrink(amount=1, axis=Axis.LAYERS)
-        assert pnp.sum(mp.mask) == mp.mask.size - size
+        assert (
+            pnp.sum(mp.full_mask(DropoutMask)) == mp.full_mask(DropoutMask).size - size
+        )
 
     def test_shrink_wire(self):
         size = 3
         mp = self._create_circuit(size)
-        mp.mask_for_axis(Axis.WIRES)[:] = True
+        mp.mask(Axis.WIRES)[:] = True
         mp.shrink(amount=1, axis=Axis.WIRES)
-        assert pnp.sum(mp.mask) == mp.mask.size - size
+        assert (
+            pnp.sum(mp.full_mask(DropoutMask)) == mp.full_mask(DropoutMask).size - size
+        )
 
     def test_shrink_parameter(self):
         size = 3
         mp = self._create_circuit(size)
-        mp.mask_for_axis(Axis.PARAMETERS)[:] = True
+        mp.mask(Axis.PARAMETERS)[:] = True
         mp.shrink(amount=1, axis=Axis.PARAMETERS)
-        assert pnp.sum(mp.mask) == mp.mask.size - 1
+        assert pnp.sum(mp.full_mask(DropoutMask)) == mp.full_mask(DropoutMask).size - 1
 
     def test_shrink_entangling(self):
         size = 3
         mp = self._create_circuit_with_entangling_gates(size)
-        mp.mask_for_axis(Axis.ENTANGLING)[:] = True
+        mp.mask(Axis.ENTANGLING)[:] = True
         mp.shrink(amount=1, axis=Axis.ENTANGLING)
-        assert (
-            pnp.sum(mp.mask_for_axis(Axis.ENTANGLING))
-            == mp.mask_for_axis(Axis.ENTANGLING).size - 1
-        )
+        assert pnp.sum(mp.mask(Axis.ENTANGLING)) == mp.mask(Axis.ENTANGLING).size - 1
 
         # also test in case no mask is set
         mp = self._create_circuit(size)
-        with pytest.raises(NotImplementedError):
+        with pytest.raises(ValueError):
             mp.shrink(amount=1, axis=Axis.ENTANGLING)
         assert Axis.ENTANGLING not in mp.masks
-        assert pnp.sum(mp.mask) == 0  # also ensure that nothing else was shrunk
+        assert (
+            pnp.sum(mp.full_mask(DropoutMask)) == 0
+        )  # also ensure that nothing else was shrunk
 
     def test_shrink_wrong_axis(self):
         mp = self._create_circuit(3)
-        with pytest.raises(NotImplementedError):
+        with pytest.raises(ValueError):
             mp.shrink(amount=1, axis=10)
 
     def test_execute(self):
@@ -267,20 +271,20 @@ class TestMaskedCircuits:
             "perturb": {
                 "amount": 1,
                 "axis": Axis.PARAMETERS,
-                "mode": Mode.ADD,
+                "mode": Mode.SET,
             }
         }
         # test empty operations
         assert MaskedCircuit.execute(mp, []) == mp
         # test existing method
         MaskedCircuit.execute(mp, [perturb_operation])
-        assert pnp.sum(mp.mask) == 1
+        assert pnp.sum(mp.full_mask(DropoutMask)) == 1
         # test existing method with copy
         new_mp = MaskedCircuit.execute(
             mp, [{"clear": {}}, {"copy": {}}, perturb_operation]
         )
         assert mp != new_mp
-        assert pnp.sum(new_mp.mask) == 1
+        assert pnp.sum(new_mp.full_mask(DropoutMask)) == 1
         # test non-existing method
         with pytest.raises(AttributeError):
             MaskedCircuit.execute(mp, [{"non_existent": {"test": 1}}])
@@ -288,28 +292,28 @@ class TestMaskedCircuits:
     def test_active(self):
         mp = self._create_circuit(3)
         assert mp.active() == 9
-        mp.mask_for_axis(Axis.WIRES)[0] = True
+        mp.mask(Axis.WIRES)[0] = True
         assert mp.active() == 6
-        mp.mask_for_axis(Axis.LAYERS)[0] = True
+        mp.mask(Axis.LAYERS)[0] = True
         assert mp.active() == 4
-        mp.mask_for_axis(Axis.PARAMETERS)[1][1] = True
+        mp.mask(Axis.PARAMETERS)[1][1] = True
         assert mp.active() == 3
 
     def test_default_value(self):
         size = 3
         mp = self._create_circuit(size)
         mp.default_value = 0
-        mp.mask_for_axis(Axis.WIRES)[0] = True
-        mp.mask_for_axis(Axis.PARAMETERS)[2, 2] = True
-        mp.mask_for_axis(Axis.LAYERS)[1] = True
+        mp.mask(Axis.WIRES)[0] = True
+        mp.mask(Axis.PARAMETERS)[2, 2] = True
+        mp.mask(Axis.LAYERS)[1] = True
         assert pnp.sum(mp.parameters[:, 0] != 0) == size
-        mp.mask_for_axis(Axis.WIRES)[0] = False
+        mp.mask(Axis.WIRES)[0] = False
         assert pnp.sum(mp.parameters[:, 0] == 0) == size
-        mp.mask_for_axis(Axis.WIRES)[1] = False
+        mp.mask(Axis.WIRES)[1] = False
         assert pnp.sum(mp.parameters[:, 1] != 0) == size
-        mp.mask_for_axis(Axis.LAYERS)[1] = False
+        mp.mask(Axis.LAYERS)[1] = False
         assert pnp.sum(mp.parameters == 0) == size * 2 - 1
-        mp.mask_for_axis(Axis.PARAMETERS)[2, 2] = False
+        mp.mask(Axis.PARAMETERS)[2, 2] = False
         assert pnp.sum(mp.parameters == 0) == size * 2
 
     def test_default_value_perturb(self):
@@ -319,7 +323,7 @@ class TestMaskedCircuits:
             wires=3,
             default_value=0,
         )
-        mp.mask_for_axis(Axis.PARAMETERS)[:] = True
+        mp.mask(Axis.PARAMETERS)[:] = True
         mp.perturb(axis=Axis.PARAMETERS, amount=0.5, mode=Mode.INVERT)
         assert pnp.sum(mp.parameters == 0) == round(0.5 * 4 * 3 * 2)
 
@@ -330,14 +334,14 @@ class TestMaskedCircuits:
             wires=3,
             default_value=0,
         )
-        mp.mask_for_axis(Axis.LAYERS)[:] = True
+        mp.mask(Axis.LAYERS)[:] = True
         mp.shrink(axis=Axis.LAYERS)
         assert pnp.sum(mp.parameters == 0) == 6
 
     def test_dynamic_parameters(self):
         size = 3
         circuit = self._create_circuit(size)
-        circuit.mask_for_axis(Axis.LAYERS)[0] = True
+        circuit.mask(Axis.LAYERS)[0] = True
         assert circuit.differentiable_parameters.size == (size - 1) * size
         assert (
             circuit.expanded_parameters(
@@ -364,13 +368,13 @@ class TestMaskedCircuits:
         rotations = [pnp.random.choice([0, 1, 2]) for _ in range(size * size)]
         circuit = qml.QNode(
             original_variational_circuit,
-            device(mp.mask_for_axis(Axis.WIRES).size),
+            device(mp.mask(Axis.WIRES).size),
         )
 
         circuit(mp.differentiable_parameters, rotations, mp)
         assert circuit.specs["gate_types"]["CZ"] == 12
 
-        mp.perturb(axis=Axis.ENTANGLING, mode=Mode.ADD, amount=6)
+        mp.perturb(axis=Axis.ENTANGLING, mode=Mode.SET, amount=6)
         circuit(mp.differentiable_parameters, rotations, mp)
         assert circuit.specs["gate_types"]["CZ"] == 6
 
@@ -386,7 +390,7 @@ class TestMaskedCircuits:
             parameters=parameters,
             layers=size,
             wires=size,
-            entangling_mask=Mask(shape=(size, size - 1)),
+            entangling_mask=DropoutMask(shape=(size, size - 1)),
         )
 
 
@@ -395,81 +399,59 @@ class TestFreezableMaskedCircuit:
         size = 3
         mp = create_freezable_circuit(size)
         assert mp
-        with pytest.raises(NotImplementedError):
-            FreezableMaskedCircuit(
-                parameters=pnp.random.uniform(low=0, high=1, size=(size, size)),
-                layers=size,
-                wires=size,
-                freeze_masks=[(Axis.ENTANGLING, Mask)],
-            )
-
-        mp = FreezableMaskedCircuit.full_circuit(
-            parameters=pnp.random.uniform(low=0, high=1, size=(size, size)),
-            layers=size,
-            wires=size,
-            wire_mask=pnp.ones((size,), dtype=bool),
-            entangling_mask=Mask(shape=(size, size - 1)),
-        )
-        assert mp
-
-    def test_mask(self):
-        size = 3
-        mp = FreezableMaskedCircuit(
-            parameters=pnp.random.uniform(low=0, high=1, size=(size, size)),
-            layers=size,
-            wires=size,
-        )
-        assert mp.mask.size == size * size
 
     def test_freeze(self):
         size = 3
         mp = create_freezable_circuit(size)
         assert (
-            mp.differentiable_parameters.size == mp.mask_for_axis(Axis.PARAMETERS).size
+            mp.differentiable_parameters.size
+            == mp.mask(Axis.PARAMETERS, mask_type=FreezeMask).size
         )
         # Test 0 amount
-        mask = mp.mask
-        mp.freeze(axis=Axis.LAYERS, amount=0, mode=Mode.ADD)
-        assert pnp.array_equal(mp.mask, mask)
+        # FIXME: does not test the same thing, befor it tested the whole mask
+        mask = mp.full_mask(FreezeMask)
+        mp.perturb(axis=Axis.LAYERS, amount=0, mode=Mode.SET, mask=FreezeMask)
+        assert pnp.array_equal(mp.full_mask(FreezeMask), mask)
         # Test freezing of layers
-        mp.freeze(axis=Axis.LAYERS, amount=1, mode=Mode.ADD)
-        assert (
-            mp.differentiable_parameters.size
-            == mp.mask_for_axis(Axis.PARAMETERS).size - size
-        )
-        assert pnp.sum(mp.freeze_mask_for_axis(Axis.LAYERS)) == 1
-        assert pnp.sum(mp.mask) == size
+        mp.perturb(axis=Axis.LAYERS, amount=1, mode=Mode.SET, mask=FreezeMask)
+        assert mp.differentiable_parameters.size == mp.mask(Axis.PARAMETERS).size - size
+        assert pnp.sum(mp.mask(axis=Axis.LAYERS, mask_type=FreezeMask)) == 1
+        assert pnp.sum(mp._accumulated_mask()) == size  # dropout and freeze mask
         # Test freezing of wires
-        mp.freeze(axis=Axis.WIRES, amount=1, mode=Mode.ADD)
-        assert pnp.sum(mp.freeze_mask_for_axis(Axis.WIRES)) == 1
-        assert pnp.sum(mp.mask) == 2 * size - 1
+        mp.perturb(axis=Axis.WIRES, amount=1, mode=Mode.SET, mask=FreezeMask)
+        assert pnp.sum(mp.mask(axis=Axis.WIRES, mask_type=FreezeMask)) == 1
+        assert (
+            pnp.sum(mp._accumulated_mask()) == 2 * size - 1
+        )  # dropout and freeze mask
         # Test freezing of parameters
-        mp.freeze(axis=Axis.PARAMETERS, amount=1, mode=Mode.ADD)
-        assert pnp.sum(mp.freeze_mask_for_axis(Axis.PARAMETERS)) == 1
-        assert pnp.sum(mp.mask) == 2 * size - 1 or pnp.sum(mp.mask) == 2 * size
+        mp.perturb(axis=Axis.PARAMETERS, amount=1, mode=Mode.SET, mask=FreezeMask)
+        assert pnp.sum(mp.mask(axis=Axis.PARAMETERS, mask_type=FreezeMask)) == 1
+        assert (
+            pnp.sum(mp._accumulated_mask()) == 2 * size - 1
+            or pnp.sum(mp._accumulated_mask()) == 2 * size  # dropout and freeze mask
+        )
         # Test wrong axis
-        with pytest.raises(NotImplementedError):
-            mp.freeze(axis=10, amount=1, mode=Mode.ADD)
+        with pytest.raises(ValueError):
+            mp.perturb(axis=10, amount=1, mode=Mode.SET, mask=FreezeMask)
 
     def test_copy(self):
         mp = create_freezable_circuit(3)
-        mp.perturb(amount=5, mode=Mode.ADD)
-        mp.freeze(amount=2, axis=Axis.LAYERS, mode=Mode.ADD)
+        mp.perturb(amount=5, mode=Mode.SET)
+        mp.perturb(amount=2, axis=Axis.LAYERS, mode=Mode.SET, mask=FreezeMask)
         mp_copy = mp.copy()
-        assert isinstance(mp_copy, FreezableMaskedCircuit)
-        assert pnp.array_equal(mp.mask, mp_copy.mask)
-        mp.perturb(amount=5, mode=Mode.REMOVE)
-        mp.freeze(amount=2, axis=Axis.LAYERS, mode=Mode.REMOVE)
-        assert pnp.sum(mp.mask) == 0
-        assert not pnp.array_equal(mp.mask, mp_copy.mask)
+        assert pnp.array_equal(mp.full_mask(FreezeMask), mp_copy.full_mask(FreezeMask))
+        mp.perturb(amount=5, mode=Mode.RESET)
+        mp.perturb(amount=2, axis=Axis.LAYERS, mode=Mode.RESET, mask=FreezeMask)
+        assert pnp.sum(mp.full_mask(FreezeMask)) == 0
+        assert not pnp.array_equal(
+            mp.full_mask(FreezeMask), mp_copy.full_mask(FreezeMask)
+        )
 
     def test_complex(self):
         random.seed(1234)
         pnp.random.seed(1234)
         mp = create_freezable_circuit(3, layer_size=2)
-        circuit = qml.QNode(
-            variational_circuit, device(mp.mask_for_axis(Axis.WIRES).size)
-        )
+        circuit = qml.QNode(variational_circuit, device(mp.mask(Axis.WIRES).size))
         optimizer = qml.GradientDescentOptimizer()
 
         def cost_fn(params, masked_circuit=None):
@@ -479,17 +461,17 @@ class TestFreezableMaskedCircuit:
                 masked_circuit,
             )
 
-        mp.freeze(axis=Axis.LAYERS, amount=2, mode=Mode.ADD)
-        mp.freeze(axis=Axis.WIRES, amount=2, mode=Mode.ADD)
+        mp.perturb(axis=Axis.LAYERS, amount=2, mode=Mode.SET, mask=FreezeMask)
+        mp.perturb(axis=Axis.WIRES, amount=2, mode=Mode.SET, mask=FreezeMask)
 
-        last_changeable = pnp.sum(mp.parameters[~mp.mask])
-        frozen = pnp.sum(mp.parameters[mp.mask])
+        last_changeable = pnp.sum(mp.parameters[~mp._accumulated_mask()])
+        frozen = pnp.sum(mp.parameters[mp._accumulated_mask()])
         for _ in range(10):
             params = optimizer.step(
                 cost_fn, mp.differentiable_parameters, masked_circuit=mp
             )
             mp.differentiable_parameters = params
-            current_changeable = pnp.sum(mp.parameters[~mp.mask])
+            current_changeable = pnp.sum(mp.parameters[~mp._accumulated_mask()])
             assert last_changeable - current_changeable != 0
-            assert frozen - pnp.sum(mp.parameters[mp.mask]) == 0
+            assert frozen - pnp.sum(mp.parameters[mp._accumulated_mask()]) == 0
             last_changeable = current_changeable
diff --git a/tests/test_masks.py b/tests/test_masks.py
index 55aec89..7932430 100644
--- a/tests/test_masks.py
+++ b/tests/test_masks.py
@@ -3,23 +3,23 @@ import pytest
 
 import pennylane.numpy as pnp
 
-from maskit._masks import Mask, PerturbationMode
+from maskit._masks import DropoutMask, ValueMask, PerturbationMode
 
 
-class TestMask:
+class TestDropoutMask:
     def test_init(self):
         size = 3
         with pytest.raises(AssertionError):
-            Mask((size,), mask=pnp.array([True, True, False, False]))
+            DropoutMask((size,), mask=pnp.array([True, True, False, False]))
         with pytest.raises(AssertionError):
-            Mask((size,), mask=pnp.array([0, 1, 3]))
+            DropoutMask((size,), mask=pnp.array([0, 1, 3]))
         preset = [False, True, False]
-        mp = Mask((size,), mask=pnp.array(preset))
+        mp = DropoutMask((size,), mask=pnp.array(preset))
         assert pnp.array_equal(mp.mask, preset)
 
     def test_setting(self):
         size = 3
-        mp = Mask((size,))
+        mp = DropoutMask((size,))
         assert mp
         assert len(mp.mask) == mp.mask.size
         assert pnp.sum(mp.mask) == 0
@@ -40,22 +40,22 @@ class TestMask:
             mp[1, 2] = True
 
     def test_wrong_mode(self):
-        mp = Mask((3,))
+        mp = DropoutMask((3,))
         with pytest.raises(NotImplementedError):
-            mp.perturb(mode=10)
+            mp.perturb(mode=10, amount=1)
 
     def test_perturbation(self):
         size = 3
-        mp = Mask((size,))
+        mp = DropoutMask((size,))
 
         for i in range(1, size + 1):
             mp.perturb(i)
-            mp.perturb(i, mode=PerturbationMode.REMOVE)
+            mp.perturb(i, mode=PerturbationMode.RESET)
             assert pnp.sum(mp.mask) == 0
 
     def test_percentage_perturbation(self):
         size = 3
-        mp = Mask((size,))
+        mp = DropoutMask((size,))
 
         for i in [0.01, 0.1, 0.5, 0.9]:
             mp.perturb(amount=i)
@@ -64,7 +64,7 @@ class TestMask:
 
     def test_wrong_percentage_perturbation(self):
         size = 3
-        mp = Mask((size,))
+        mp = DropoutMask((size,))
 
         for i in [1.1, 1.5, 3.1]:
             mp.perturb(amount=i)
@@ -72,51 +72,51 @@ class TestMask:
             mp.clear()
 
     def test_negative_perturbation(self):
-        mp = Mask((3,))
+        mp = DropoutMask((3,))
         with pytest.raises(AssertionError):
             mp.perturb(amount=-1)
 
     def test_perturbation_remove_add(self):
         size = 3
-        mp = Mask((size,))
+        mp = DropoutMask((size,))
 
         for amount in [random.randrange(size), 0, size, size + 1]:
-            mp.perturb(amount=amount, mode=PerturbationMode.REMOVE)
+            mp.perturb(amount=amount, mode=PerturbationMode.RESET)
             assert pnp.sum(mp.mask) == 0
-            mp.perturb(amount=amount, mode=PerturbationMode.ADD)
+            mp.perturb(amount=amount, mode=PerturbationMode.SET)
             assert pnp.sum(mp.mask) == min(amount, size)
             mp.clear()
 
     def test_perturbation_invert_remove(self):
         size = 3
-        mp = Mask((size,))
+        mp = DropoutMask((size,))
 
         for amount in [random.randrange(size), 0, size, size + 1]:
             mp.perturb(amount=amount, mode=PerturbationMode.INVERT)
             reversed_amount = pnp.sum(mp.mask).unwrap()  # unwrap tensor
-            mp.perturb(amount=reversed_amount, mode=PerturbationMode.REMOVE)
+            mp.perturb(amount=reversed_amount, mode=PerturbationMode.RESET)
             assert pnp.sum(mp.mask) == 0
 
     def test_perturbation_add_remove(self):
         size = 3
-        mp = Mask((size,))
+        mp = DropoutMask((size,))
 
         for amount in [random.randrange(size), 0, size, size + 1]:
-            mp.perturb(amount=amount, mode=PerturbationMode.ADD)
+            mp.perturb(amount=amount, mode=PerturbationMode.SET)
             assert pnp.sum(mp.mask) == min(amount, size)
-            mp.perturb(amount=amount, mode=PerturbationMode.REMOVE)
+            mp.perturb(amount=amount, mode=PerturbationMode.RESET)
             assert pnp.sum(mp.mask) == 0
 
     @pytest.mark.parametrize(
         "mode",
         [
-            (PerturbationMode.ADD, PerturbationMode.INVERT),
+            (PerturbationMode.SET, PerturbationMode.INVERT),
             (PerturbationMode.INVERT, PerturbationMode.INVERT),
         ],
     )
     def test_perturbation_mode(self, mode):
         size = 3
-        mp = Mask((size,))
+        mp = DropoutMask((size,))
 
         for amount in [0, size, size + 1]:
             mp.perturb(amount=amount, mode=mode[0])
@@ -125,7 +125,7 @@ class TestMask:
 
     def test_shrink(self):
         size = 3
-        mp = Mask((size,))
+        mp = DropoutMask((size,))
 
         for amount in range(size + 1):
             mp[:] = True
@@ -134,7 +134,7 @@ class TestMask:
 
     def test_shrink_nd(self):
         size = 3
-        mp = Mask((size, size - 1))
+        mp = DropoutMask((size, size - 1))
         for amount in range(mp.mask.size + 1):
             mp[:] = True
             mp.shrink(amount)
@@ -142,7 +142,7 @@ class TestMask:
 
     def test_copy(self):
         size = 3
-        mp = Mask((size,))
+        mp = DropoutMask((size,))
         new_mp = mp.copy()
         mp[0] = True
         assert pnp.sum(mp.mask) > pnp.sum(new_mp.mask)
@@ -150,10 +150,43 @@ class TestMask:
 
     def test_apply_mask(self):
         size = 3
-        mp = Mask((size,))
+        mp = DropoutMask((size,))
         with pytest.raises(IndexError):
             mp.apply_mask(pnp.ones((size - 1, size)))
         mp.mask[1] = True
         result = mp.apply_mask(pnp.ones((size,), dtype=bool))
         assert pnp.sum(mp.mask) == 1
         assert pnp.sum(result) == size - 1
+
+
+class TestValueMask:
+    def test_apply_mask(self):
+        size = 3
+        mp = ValueMask((size,))
+        with pytest.raises(ValueError):
+            mp.apply_mask(pnp.ones((size - 1,)))
+        mp.mask[1] = 1
+        result = mp.apply_mask(pnp.ones((size,), dtype=float))
+        assert pnp.sum(mp.mask) == 1
+        assert pnp.sum(result) == size + 1
+        mp.mask[1] = -1
+        result = mp.apply_mask(pnp.ones((size,), dtype=float))
+        assert pnp.sum(result) == size - 1
+
+    def test_perturbation(self):
+        size = 3
+        mp = ValueMask((size,))
+
+        for i in range(1, size + 1):
+            mp.perturb(i, mode=PerturbationMode.SET, value=1)
+            mp.perturb(i, mode=PerturbationMode.RESET)
+            assert pnp.sum(mp.mask) == 0
+
+    def test_shrink(self):
+        size = 3
+        mp = ValueMask((size,))
+
+        for amount in range(size + 1):
+            mp[:] = 1
+            mp.shrink(amount)
+            assert pnp.sum(mp.mask) == size - amount
diff --git a/tests/test_optimizers.py b/tests/test_optimizers.py
index 9aaf472..de5b4fa 100644
--- a/tests/test_optimizers.py
+++ b/tests/test_optimizers.py
@@ -27,7 +27,7 @@ class TestLBFGSBOptimizer:
         mp_step = mp.copy()
         circuit = qml.QNode(
             plain_variational_circuit,
-            device(mp.mask_for_axis(Axis.WIRES).size),
+            device(mp.mask(Axis.WIRES).size),
         )
 
         def cost_fn(params):
@@ -57,7 +57,7 @@ class TestLBFGSBOptimizer:
         mp_original = mp.copy()
         circuit = qml.QNode(
             plain_variational_circuit,
-            device(mp.mask_for_axis(Axis.WIRES).size),
+            device(mp.mask(Axis.WIRES).size),
         )
 
         def cost_fn(params):
@@ -77,9 +77,7 @@ class TestLBFGSBOptimizer:
         optimizer = L_BFGS_B()
         mp = create_circuit(3, layer_size=2)
         mp_step = mp.copy()
-        circuit = qml.QNode(
-            variational_circuit, device(mp.mask_for_axis(Axis.WIRES).size)
-        )
+        circuit = qml.QNode(variational_circuit, device(mp.mask(Axis.WIRES).size))
 
         def cost_fn(params, masked_circuit=None):
             return cost(
@@ -105,9 +103,7 @@ class TestLBFGSBOptimizer:
         optimizer = L_BFGS_B()
         ensemble = Ensemble(dropout=RANDOM)
         mp = create_circuit(3, layer_size=2)
-        circuit = qml.QNode(
-            variational_circuit, device(mp.mask_for_axis(Axis.WIRES).size)
-        )
+        circuit = qml.QNode(variational_circuit, device(mp.mask(Axis.WIRES).size))
 
         def cost_fn(params, masked_circuit=None):
             return cost(
diff --git a/tests/utils.py b/tests/utils.py
index 2277e2d..4377643 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -1,8 +1,8 @@
 import pennylane as qml
 import pennylane.numpy as pnp
 
-from maskit._masks import PerturbationAxis as Axis
-from maskit._masked_circuits import FreezableMaskedCircuit, MaskedCircuit
+from maskit._masks import PerturbationAxis as Axis, DropoutMask, FreezeMask
+from maskit._masked_circuits import MaskedCircuit
 
 
 def device(wires: int):
@@ -34,32 +34,38 @@ def create_freezable_circuit(size: int, layer_size: int = 1):
         parameters = pnp.random.uniform(
             low=-pnp.pi, high=pnp.pi, size=(size, size, layer_size)
         )
-    return FreezableMaskedCircuit.full_circuit(
-        parameters=parameters, layers=size, wires=size
+    return MaskedCircuit(
+        parameters=parameters,
+        layers=size,
+        wires=size,
+        masks=(
+            (axis, mask_type)
+            for axis in Axis
+            if axis is not Axis.ENTANGLING
+            for mask_type in [DropoutMask, FreezeMask]
+        ),
     )
+    # TODO: remove me
+    # return FreezableMaskedCircuit.full_circuit(
+    #     parameters=parameters, layers=size, wires=size
+    # )
 
 
 def variational_circuit(params, masked_circuit: MaskedCircuit = None):
     full_parameters = masked_circuit.expanded_parameters(params)
-    for layer, layer_hidden in enumerate(masked_circuit.mask_for_axis(Axis.LAYERS)):
+    for layer, layer_hidden in enumerate(masked_circuit.mask(Axis.LAYERS)):
         if not layer_hidden:
-            for wire, wire_hidden in enumerate(
-                masked_circuit.mask_for_axis(Axis.WIRES)
-            ):
+            for wire, wire_hidden in enumerate(masked_circuit.mask(Axis.WIRES)):
                 if not wire_hidden:
-                    if not masked_circuit.mask_for_axis(Axis.PARAMETERS)[layer][wire][
-                        0
-                    ]:
+                    if not masked_circuit.mask(Axis.PARAMETERS)[layer][wire][0]:
                         qml.RX(full_parameters[layer][wire][0], wires=wire)
-                    if not masked_circuit.mask_for_axis(Axis.PARAMETERS)[layer][wire][
-                        1
-                    ]:
+                    if not masked_circuit.mask(Axis.PARAMETERS)[layer][wire][1]:
                         qml.RY(full_parameters[layer][wire][1], wires=wire)
-            for wire in range(0, masked_circuit.mask_for_axis(Axis.LAYERS).size - 1, 2):
+            for wire in range(0, masked_circuit.mask(Axis.LAYERS).size - 1, 2):
                 qml.CZ(wires=[wire, wire + 1])
-            for wire in range(1, masked_circuit.mask_for_axis(Axis.LAYERS).size - 1, 2):
+            for wire in range(1, masked_circuit.mask(Axis.LAYERS).size - 1, 2):
                 qml.CZ(wires=[wire, wire + 1])
-    return qml.probs(wires=range(len(masked_circuit.mask_for_axis(Axis.WIRES))))
+    return qml.probs(wires=range(len(masked_circuit.mask(Axis.WIRES))))
 
 
 def plain_variational_circuit(params):