def assert_listens_to_forceful_enabled(self, layer_constructor):
        """
        Setup to test that the stochastic mode is read correctly in a noise layer in eager mode.
        To be called directly by the test class for the corresponding layer

        Args:
            layer_constructor (): No-Args lambda that creates a corresponding layer.

        Returns:
            None
        """
        stochastic_mode = StochasticMode()
        self.assertTrue(
            tf.executing_eagerly(), "This test is supposed to test eager execution"
        )
        layer = layer_constructor(stochastic_mode)
        stochastic_mode.as_tensor().assign(tf.ones((), dtype=bool))
        enabled_result = layer(DUMMY_INPUT)
        self.assertFalse(
            tf.reduce_all(tf.equal(enabled_result, DUMMY_INPUT).numpy()),
            "No values were dropped",
        )
        stochastic_mode.as_tensor().assign(tf.zeros((), dtype=bool))
        disabled_result = layer(DUMMY_INPUT)
        self.assertTrue(
            tf.reduce_all(tf.equal(disabled_result, DUMMY_INPUT)).numpy(),
            "Some values changed - which should not happen",
        )
 def test_cast_from_keras(self):
     plain_keras = tf.keras.layers.GaussianNoise(stddev=0.5)
     stochastic_mode = StochasticMode()
     output = UwizGaussianNoise.from_keras_layer(
         plain_keras, stochastic_mode=stochastic_mode
     )
     self.assertIsInstance(output, UwizGaussianNoise)
     self.assertEqual(output.stddev, 0.5)
     self.assertEqual(
         stochastic_mode.as_tensor(), output.stochastic_mode.as_tensor()
     )
 def test_cast_from_keras(self):
     plain_keras = tf.keras.layers.Dropout(rate=0.5)
     stochastic_mode = StochasticMode()
     output = UwizBernoulliDropout.from_keras_layer(
         plain_keras, stochastic_mode=stochastic_mode
     )
     self.assertIsInstance(output, UwizBernoulliDropout)
     self.assertEqual(output.rate, 0.5)
     self.assertEqual(
         stochastic_mode.as_tensor(), output.stochastic_mode.as_tensor()
     )
    def __init__(self,
                 inputs,
                 outputs,
                 stochastic_mode: StochasticMode,
                 name: str = None):
        """
        Create a new functional model, equivalent to calling tf.keras.Model(...).

        In addition, a stochastic mode instance has to be passed.
        The same instance also has to be passed to any randomized uwiz.layers instances
        which are part of this model.
        This allows to dynamically enable and disable randomness in the predictions.
        :param inputs: See the corresponding tf.keras.Model(...) docs
        :param outputs: See the corresponding tf.keras.Model(...) docs
        :param stochastic_mode: A stochastic mode instance
        :param name: See the corresponding tf.keras.Model(...) docs
        """
        super().__init__()
        self._inner_model = tf.keras.Model(inputs=inputs,
                                           outputs=outputs,
                                           name=name)
        self._inner_model._stochastic_mode_tensor = stochastic_mode.as_tensor()
def stochastic_from_keras(
    model: tf.keras.models.Model,
    input_tensors=None,
    clone_function=None,
    expect_determinism=False,
    temp_weights_path="tmp/weights",
):
    """
    Creates a stochastic instance from a given `tf.keras.models.Sequential` model:
    The new model will have the same structure (layers) and weights as the passed model.

    All stochastic layers (e.g. tf.keras.layers.Dropout) will be used for randomization during randomized predictions.
    If no stochastic layers are present, a ValueError is thrown.
    The raising of the error can be suppressed by setting `expect_determinism` to true.

    If your model contains custom layers, you can pass a function to `clone_function` to clone your custom layers,
    or place the annotation `@tf.keras.utils.register_keras_serializable()` on your custom layers,
    and make sure the `get_config` and `from_config` methods are implemented.
    (uncertainty wizard will serialize and deserialize all layers).

    :param model: The model to copy. Remains unchanged.
    :param input_tensors: Optional tensors to use as input_tensors for new model. See the corresponding parameter in `tf.keras.models.clone_model` for details.
    :param _clone_function: Optional function to use to clone layers. Will be applied to all layers except input layers and stochastic layers. See the corresponding parameter in `tf.keras.models.clone_model` for more details.
    :param expect_determinism: If True, deterministic models (e.g. models without stochastic layers) are accepted and no ValueError is thrown.
    :param temp_weights_path: The model weights are temporarily saved to the disk at this path. Folder is deleted after successful completion.
    :return: A newly created stochastic model
    """
    # _clone_function is some layer cloning behavior that can be specified by the user
    # If none is specified, we use keras default (see `tf.keras.models.clone_model`)
    if clone_function is None:

        def _clone_function(layer):
            return layer.__class__.from_config(layer.get_config())

    # We wrap the users (or default) clone function in a clone function
    #   that replaces stochastic layers with uncertainty wizard stochastic layers
    stochastic_mode = StochasticMode()
    is_stochastic_layer = []

    def _uncertainty_wizard_aware_clone_function(layer):
        new_layer = Stochastic._replace_layer_if_possible(
            layer, stochastic_mode=stochastic_mode
        )
        if new_layer == layer:
            # Layer was not mapped to an uncertainty wizard layer, thus the default clone function is applied
            new_layer = _clone_function(layer)
            is_stochastic_layer.append(False)
        else:
            is_stochastic_layer.append(True)
        return new_layer

    # Clone the keras model to become the new inner model
    new_inner = tf.keras.models.clone_model(
        model=model,
        input_tensors=input_tensors,
        clone_function=_uncertainty_wizard_aware_clone_function,
    )
    new_inner.stochastic_mode_tensor = stochastic_mode.as_tensor()

    if not expect_determinism and not any(is_stochastic_layer):
        raise ValueError(
            "The passed model had no stochastic layers."
            "If that is intended (and you do not plan to use any sampling based quantifiers)"
            "you can set the flag `expect_determinism = True`, i.e., "
            "calling `SequentialStochastic.clone_from_keras(keras_model,expect_determinism = True)`"
        )

    # Restore the Weights
    model.save_weights(temp_weights_path)
    new_inner.load_weights(temp_weights_path)
    # Remove temporarily stored weights
    shutil.rmtree(temp_weights_path, ignore_errors=True)

    # Put the wrapper around the new model
    if isinstance(model, tf.keras.models.Sequential):
        target_class = StochasticSequential
    else:
        target_class = StochasticFunctional

    # Consenting Adults: The _wrap method is intended to be used here
    #   but declared private as it is not intended to be used by the uwiz user
    return target_class._wrap(
        inner=new_inner, stochastic_mode_tensor=stochastic_mode.as_tensor()
    )