vectorized_map.py

import tensorflow as tf
tf.enable_v2_behavior()

def outer_product(x):
  return tf.tensordot(x, x, 0)

val = tf.ones((1000, 32, 32))
%timeit tf.map_fn(outer_product, val)
# >>> 1 loops, best of 3: 849 ms per loop
%timeit tf.vectorized_map(outer_product, val)
# >>> 10 loops, best of 3: 37.8 ms per loop