I have also been troubled by this problem recently, I need to process the output of the neural network, so I need to call a function written in pure numpy to calculate the loss after the output is processed. When using tf. py_function in tensorflow, I found that for functions not operated by tf, although py_function can get the calculation result, this result cannot be used to save gradient for backpropagation.
tf.py_function(func=external_func, inp=[input_external_func], Tout=tf.float32)
There should be no solution to this problem at present. The external functions I need to call are complex FEM simulation libraries that I can't implement from scratch with tensorflow or pytorch.
reference resourse:
How to use a numpy-based external library function within a Tensorflow 2/keras custom layer?