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, layers=None, name=None):
     super().__init__()
     # Create an empty sequential model with a stochastic mode
     self._inner_sequential = tf.keras.models.Sequential(name=name)
     self._inner_sequential._stochastic_mode_tensor = StochasticMode().as_tensor()
     # Append all the passed layers
     if layers is not None:
         for layer in layers:
             self.add(layer)
 def _wrap(cls, inner: tf.keras.Model, stochastic_mode_tensor=None):
     if stochastic_mode_tensor is None:
         assert inner._stochastic_mode_tensor is not None, (
             "Uncertainty Wizard internal error. "
             "Trying to wrap a model that has no stochastic_mode_tensor, "
             "and no external stochastic_mode_tensor is passed to attach")
         stochastic_mode_tensor = inner._stochastic_mode_tensor
     stochastic_mode = StochasticMode(stochastic_mode_tensor)
     return StochasticFunctional(inner.inputs,
                                 inner.outputs,
                                 stochastic_mode=stochastic_mode)
    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 assert_listens_to_forceful_enabled_graph_mode(self, layer_constructor):
        """
        Setup to test that the stochastic mode is read correctly in a noise layer.
        To be called directly by the test class for the corresponding layer

        Args:
            layer_constructor (): function that creates a layer. Argument: the stochastic mode instance to use

        Returns:
            None
        """
        stochastic_mode = StochasticMode()

        @tf.function
        def run_test():
            self.assertFalse(
                tf.executing_eagerly(),
                "This test is supposed to test disabled eager execution",
            )
            layer = layer_constructor(stochastic_mode)
            stochastic_mode.as_tensor().assign(tf.ones((), dtype=bool))
            enabled_result = layer(DUMMY_INPUT)
            input_as_tensor = tf.constant(DUMMY_INPUT, dtype=tf.float32)
            test_enabled = tf.debugging.Assert(
                tf.math.logical_not(
                    tf.reduce_all(tf.equal(enabled_result, input_as_tensor))
                ),
                [tf.constant("No values were dropped"), enabled_result],
            )
            stochastic_mode.as_tensor().assign(tf.zeros((), dtype=bool))
            disabled_result = layer(DUMMY_INPUT)
            test_disabled = tf.debugging.Assert(
                tf.reduce_all(tf.equal(disabled_result, input_as_tensor)),
                [
                    tf.constant("Some values changed - which should not happen"),
                    disabled_result,
                ],
            )
            stochastic_mode.as_tensor().assign(tf.zeros((), dtype=bool))

            with tf.control_dependencies([test_enabled, test_disabled]):
                pass

        run_test()
    def test_warns_custom_uwiz_subtype(self):
        class SubUwizGaussianNoise(UwizGaussianNoise):
            def __init__(self, stddev, stochastic_mode, **kwargs):
                super().__init__(stddev, stochastic_mode, **kwargs)

        subclass_instance = SubUwizGaussianNoise(
            stddev=0.5, stochastic_mode=StochasticMode()
        )
        uwiz_type = uwiz.models.stochastic_utils.layers.UwizGaussianNoise
        with self.assertWarns(UncertaintyWizardWarning) as cm:
            uwiz.models.stochastic_utils.layers._has_casting_preventing_subtype(
                subclass_instance,
                expected_type=tf.keras.layers.GaussianNoise,
                corresponding_uw_type=uwiz_type,
            )
        the_warning = cm.warning
        self.assertTrue(
            "ou know what you did and set up the stochastic mode correctly."
            in the_warning.args[0]
        )
    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 _get_stochastic_mode(self):
     """
     Get access to the stochastic mode used in the model to toggle randomness during predictions.
     :return: A stochastic mode instance wrapping the stochastic mode tensor used in this model.
     """
     return StochasticMode(self.stochastic_mode_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()
    )