def test_tf2_mnist_certification(art_warning, fix_get_mnist_data):
    """
    Assert that the correct number of certifications are given for the MNIST dataset
    """

    import tensorflow as tf

    def build_model(input_shape):
        img_inputs = tf.keras.Input(shape=(input_shape[0], input_shape[1], input_shape[2] * 2))
        x = tf.keras.layers.Conv2D(filters=32, kernel_size=(4, 4), strides=(2, 2), activation="relu")(img_inputs)
        x = tf.keras.layers.MaxPool2D(pool_size=(2, 2), strides=2)(x)
        # tensorflow uses channels last, and we are loading weights from an originally trained pytorch model
        x = tf.transpose(x, (0, 3, 1, 2))
        x = tf.keras.layers.Flatten()(x)
        x = tf.keras.layers.Dense(100, activation="relu")(x)
        x = tf.keras.layers.Dense(10)(x)
        return tf.keras.Model(inputs=img_inputs, outputs=x)

    def get_weights():
        fpath = os.path.join(
            os.path.dirname(os.path.dirname(__file__)), "../../utils/resources/models/certification/derandomized/"
        )
        weight_names = [
            "W_CONV2D1_MNIST.npy",
            "B_CONV2D1_MNIST.npy",
            "W_DENSE1_MNIST.npy",
            "B_DENSE1_MNIST.npy",
            "W_DENSE2_MNIST.npy",
            "B_DENSE2_MNIST.npy",
        ]
        weight_list = []
        for name in weight_names:
            w = np.load(fpath + name)
            if "W_CONV" in name:
                w = np.transpose(w, (2, 3, 1, 0))
            if "W_DENSE" in name:
                w = np.transpose(w)
            weight_list.append(w)
        return weight_list

    net = build_model(input_shape=(28, 28, 1))
    net.set_weights(get_weights())

    loss_object = tf.keras.losses.CategoricalCrossentropy(from_logits=True)
    optimizer = tf.keras.optimizers.legacy.SGD(learning_rate=0.01)

    try:
        for ablation_type in ["column", "block"]:
            if ablation_type == "column":
                size_to_certify = 5
                ablation_size = 2
            else:
                """
                the model was trained on column ablations, so make the block task simpler so that a
                degree of certification is obtained.
                """
                size_to_certify = 1
                ablation_size = 5

            classifier = TensorFlowV2DeRandomizedSmoothing(
                model=net,
                clip_values=(0, 1),
                loss_object=loss_object,
                optimizer=optimizer,
                input_shape=(28, 28, 1),
                nb_classes=10,
                ablation_type=ablation_type,
                ablation_size=ablation_size,
                threshold=0.3,
                logits=True,
            )

            x = np.copy(fix_get_mnist_data[0])
            x = np.squeeze(x)
            x = np.expand_dims(x, axis=-1)
            preds = classifier.predict(x)
            cert, cert_and_correct, top_predicted_class_argmax = classifier.ablator.certify(
                preds, label=fix_get_mnist_data[1], size_to_certify=size_to_certify
            )

            if ablation_type == "column":
                assert np.sum(cert) == 52
            else:
                assert np.sum(cert) == 22

    except ARTTestException as e:
        art_warning(e)

    Assert that the correct number of certifications are given for the MNIST dataset
    
tests/estimators/certification/test_derandomized_smoothing.py