Apply vmap jnp.pad over the batch in REFLECT convolution mode

[sarlinpe]

What does this PR do?

This PR fixes an error occurring when exporting flax models with jax2tf, polymorphic batch size, and REFLECT padding mode in convolutions. This issue can be minimally reproduced as

import jax.numpy as jnp
import tensorflow as tf
from jax.experimental import jax2tf

def fn_jax(x):
  return jnp.pad(x, pad_width=[(0, 0), (1, 1), (1, 1)], mode='reflect')

fn_tf = jax2tf.convert(fn_jax, polymorphic_shapes=["(b, 14, 14)"])
fn_tf = tf.function(fn_tf, autograph=False, jit_compile=True)
fn_tf.get_concrete_function(tf.TensorSpec([None, 14, 14], tf.float32))

which results in

ValueError: Shape polymorphism is supported for jnp.pad with 'reflect' or 'symmetric' padding mode only when it is possible to determine at lowering time that the axis size (= b) is larger than 1 and larger or equal than the padding length (= 0). Error while handling left padding on axis 0.

Applying jax.vmap(functools.partial(jnp.pad, pad_width=[(1, 1), (1, 1)], mode='reflect'))(x) fixes the issue.

Checklist

• This PR fixes a minor issue (e.g.: typo or small bug) or improves the docs (you can dismiss the other
  checks if that's the case).
• This change is discussed in a Github issue/
  discussion (please add a
  link).
• The documentation and docstrings adhere to the
  documentation guidelines.
• This change includes necessary high-coverage tests.
  (No quality testing = no merge!)

[sarlinpe]

@cgarciae It would be great to receive your review, thank you!