possibily best version

Author: scarlehoff

madgraph_plugin/template_files/matrix_method_python.inc

@@ -101,8 +101,10 @@ class Matrix_%(process_string)s(object):
         def loopable_matrix(hel):
             return self.matrix(all_ps, hel, %(params)s)
 
-        ans = tf.map_fn(loopable_matrix, self.helicities, back_prop=False, fn_output_signature=DTYPE)
-        return tf.reduce_sum(ans, axis=0)/self.denominator
+        ans = tf.zeros_like(all_ps[:,0,0])
+        for hel in self.helicities:
+            ans += loopable_matrix(hel)
+        return ans/self.denominator
 
 
     @tf.function(input_signature=matrix_signature)