How to use a numpy-based external library function within a Tensorflow 2/keras custom layer?
I would be very grateful if someone would help me understanding the following problem.
I am trying to implement a custom layer in tensorflow 2 by using keras (it's a layer derived from the class Layer).
I have overriden the build and the call functions.
When I write my call function I need to invoke a method from an external library which only accepts numpy array. It's a pretty complex function and of course it does not use tensorflow functions.
My call function takes in input a tensor, convert it to numpy (by .numpy() function), invoke the external method from another library, then converts the numpy array back to tensor.
Here I have some issues and tried different solutions.

If I run the code it tells me that ...ops.Tensor has no method .numpy(). If I correctly understood, this is due to the type of tensor I receive in input to the call function and I need to use the eager mode.
If I compile by setting the run_eagerly flag to true, during the fit, it tells me that the gradient is missing. I was hoping that some form of auto-grad computation was implemented but maybe this is not what it is currently happening?.

Source code:
class SimpleLayer(tf.keras.layers.Layer):
    def __init__(self):
        super(SimpleLayer, self).__init__()

    def build(self, input_shape):
        shape_a = input_shape[0]                #batch_size
        shape_b = int(input_shape[1])           #height
        shape_c = int(input_shape[2])           #width
        shape_d = input_shape[3]                #channel_number
        self.w = self.add_weight(shape=(shape_d*shape_a, 2, shape_b, shape_c),
                                 initializer='random_uniform', dtype='float64',
                                 trainable=True)


    def call(self, inputs):
        tf_inputs = inputs.numpy()
        sft = tr.transform(tf_inputs) # call to the external library
        transformed = tf.math.multiply(sft, self.w)
        conv = tf.zeros(np.shape(inputs))
        conv = tr.inv_transform(inputs, transformed) # another call to a function that also converts to tensor
        return conv. # this is a tensor!


.....

model = tf.keras.Sequential([
    layers.experimental.preprocessing.Rescaling(1./255),
    pl.SimpleLayer(),
    layers.Activation('relu'),
    layers.MaxPooling2D(),
    pl.SimpleLayer(),
    layers.Activation('relu'),
    layers.Flatten(),
    layers.Dense(128, activation='relu'),
    layers.Dense(num_classes)
])

model.compile(
    optimizer='adam',
    loss=tf.losses.SparseCategoricalCrossentropy(from_logits=True),
    run_eagerly=True,  # adding this to avoid graph building
    metrics=['accuracy'])

history = model.fit(train_images, train_labels, epochs=3, batch_size=20,
                    validation_data=(test_images, test_labels))

Is there any workaround for that? Can I only use tensorflow 2 functions to process the tensor or maybe I am missing something?
I am also considering to switch to pyTorch (which I do not know yet) but I fear to have the same issue. What would you suggest?

Gradient cannot flow through numpy functions as... numpy has no autograd capabilities of any sort.
Your custom function will not work in either tensorflow2.x nor in pytorch as the operation on numpy arrays are not recorded.
For your call code (in PyTorch it would be forward, otherwise really similar), see comments:
def call(self, inputs):
    tf_inputs = inputs.numpy() # gradient breaks
    sft = tr.transform(tf_inputs) # still numpy
    transformed = tf.math.multiply(sft, self.w)
    conv = tf.zeros(np.shape(inputs))
    conv = tr.inv_transform(inputs, transformed)
    return conv # Tensor with no gradient history

And this would happen for all your SimpleLayer layers.
Possible solutions

Code those operations in differentiable manner (assuming they are differentiable and do not use operations like argmax) either in tensorflow or pytorch (the latter is easier and more intuitive and you would have easier time with it IMO).
If your complicated functions use numpy and are differentiable then you could code it with PyTorch by rewriting np.op to torch.op as their APIs are pretty similar and interoperate almost seamlessly
You may go with JAX instead of tf and pytorch as &quot;JAX can automatically differentiate native Python and NumPy functions&quot;. Depending on your application it might be worth a shot

