[BUG] Custom keras layer no relevances

[mattfedrau]

Describe the bug

I´m trying to use a custom keras Layer(RBF) with the analyzer 'deep_taylor'(see attached code).
The analysis is 0 everywhere, so my question is whether there is no support for custom keras layer
or whether I do something wrong?
My expectation is, that there are at least a few non-zero relevances.

from keras.layers import Layer
from keras import backend as K
import keras
import tensorflow as tf
import innvestigate
import matplotlib.pyplot as plt


tf.compat.v1.disable_eager_execution()


class RBFLayer(Layer):
    def __init__(self, units, gamma, **kwargs):
        super(RBFLayer, self).__init__(**kwargs)
        self.units = units
        self.gamma = K.cast_to_floatx(gamma)

    def build(self, input_shape):

        self.mu = self.add_weight(name='mu',
                                  shape=(int(input_shape[1]), self.units),
                                  initializer='uniform',
                                  trainable=True)
        super(RBFLayer, self).build(input_shape)

    def call(self, inputs):
        diff = K.expand_dims(inputs) - self.mu
        l2 = K.sum(K.pow(diff, 2), axis=1)
        res = K.exp(-1 * self.gamma * l2)
        return res

    def compute_output_shape(self, input_shape):
        return (input_shape[0], self.units)


import pandas as pd
import numpy as np

(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

from keras.layers import Dense, Flatten
from keras.models import Sequential
from keras.losses import binary_crossentropy

model = Sequential()
model.add(Flatten(input_shape=(28, 28)))
model.add(RBFLayer(10, 0.04478))
model.add(Dense(10, activation='softmax', name='foo'))

model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])


model.summary()



y_train = tf.keras.utils.to_categorical(y_train, num_classes=10, dtype=int)
model.fit(x_train, y_train, batch_size=256, epochs=3)
model_wo_softmax = innvestigate.model_wo_softmax(model)
analyzer = innvestigate.create_analyzer('deep_taylor', model_wo_softmax, neuron_selection_mode="index")
analysis = analyzer.analyze(x_train[0].reshape(1, 28, 28), neuron_selection=4)


np.max(analysis)
plt.imshow(x_train[0])
plt.imshow(analysis[0])

Platform information

• OS: Ubuntu 22.04
• Python version: 3.10.12
• iNNvestigate version: 2.1.0
• TensorFlow version: 2.12.1
• keras 2.12.0