how to use different convolution layers in tf.map_fn?

Question

I tried to establish a simple multi-head attention layer in tensorflow1.14. Each head contains three different conv1d layers. And I want to use tf.map_fn to com