vmapping pad with variable argument but static output size

[stefano-1981]

Hi there,

I have a tensor of images, say, BxNxN and a translation tensor Bx2. I would like to vmap the jax.numpy.pad operation along the batch axis 0 like this:

def mypad(m, translation)
  out_shape = (1000, 1000)
  return jax.numpy.pad(m, ((translation[0], 1000 - N - translation[0]), (translation[1], 1000 - N - translation[1]) ))

# images has shape [B,N,N] and translations has shape [B, 2]
out = jax.vmap(mypad)(images, translations)

Now I think I understand that this does not work because the pad operation output is non static. However this is a special case where the padding is variable but the output shape of mypad() is always [1000, 1000]. Is there a way I can enforce that the output shape is indeed constant?

As a workaround I've been using jax.image.scale_and_translate but it's probably much more expensive (and besides it gives me troubles when converting the model to tflite).

Any help?

[gnecula]

Can you use dynamic_update_slice instead of pad? Something like this:

def my_pad(m, translation):
      base = jnp.zeros((1000, 1000), dtype=jnp.float32)
      return lax.dynamic_update_slice(base, m, (translation[0], translation[1]))

out = jax.vmap(my_pad)(images, translation)

If this works for you in JAX, then the next step is to see if this can be converted with jax2tf + tflite.

[stefano-1981]

Thanks, that works indeed. I have other troubles with jax2tf + tflite, but that's indeed a separate problem.