How do I use a function to compare feature vectors in a simaes twin network in tf/keras

See more:

Python

tensorflow

Keras

Hi I am quite unexperienced with that topic.
I try to compare two images and interpret this difference. My Idea is to train a siamese CNN and a standard neural network (NN). The later one should analyse the cross correlation matrix of the two feature vectors. In order to produce this matrix the following function is used:

def NCC(a, b):
	l=a.shape[1]
	av_a=tf.math.reduce_mean(a)
	av_b=tf.math.reduce_mean(b)
	a=a-av_a
	b=b-av_b
	norm_a=tf.math.sqrt(tf.math.reduce_sum(a*a))
	norm_b=tf.math.sqrt(tf.math.reduce_sum(b*b))
	a=a/norm_a
	b=b/norm_b
	A=tf.reshape(tf.repeat(a, axis=0, repeats=l),(l,l))
	B=tf.reshape(tf.repeat(b, axis=0, repeats=l),(l,l))
	ncc=Flatten()(A*tf.transpose(B))		
	return ncc

(This function produces a matrix with entries a_i*b_j for all combinations i and j)
I call NCC in the model between the CNN and the NN simply by using :

input_nn=NCC(feature_vectorA,feature_vectorB)

This function works well with batch size 1. However I get problems if I increase the batch size, due to

tf.reshape(tf.repeat(a, axis=0, repeats=l),(l,l))

because the input vector has the shape (batch_size,l). I need for the NN (after NCC) the output vector with shape (batch_size,l^2). Any hint how I can manage this, is appreciated.

What I have tried:

I tried also to use this function in a Lamda-Layer, e.g.

input_nn=Lambda(NCC)(feature_vectorA,feature_vectorB)

However this seems not to work