tfp.experimental.nn.util.convolution_batch

Like tf.nn.conv2d except applies batch of kernels to batch of x.