Inverse operation to padding in Jax
I'm trying to learn how to use Jax and I stumbled upon the problem of converting the torch.nn.functionnal.pad function into Jax. There is a function to perform padding but I would like in the same way as in PyTorch use negative numbers in the padding (e.g F.pad(array, [-1,-1])).
Does anyone have an idea or had the same problem ?

The jax.lax.pad function accepts negative padding indices, although the API is a bit different than that of torch.nn.functional.pad. For example:
from jax import lax
import jax.numpy as jnp

x = jnp.ones((2, 3))
y = lax.pad(x, padding_config=[(0, 0, 0), (1, 1, 0)], padding_value=0.0)
print(y)
# [[0. 1. 1. 1. 0.]
#  [0. 1. 1. 1. 0.]]

x = lax.pad(y, padding_config=[(0, 0, 0), (-1, -1, 0)], padding_value=0.0)
print(x)
# [[1. 1. 1.]
#  [1. 1. 1.]]

If you wish, you could wrap this with a function that has similar semantics to the torch version. Here's a quick attempt:
def jax_pad(input, pad, mode='constant', value=0):
  &quot;&quot;&quot;JAX implementation of torch.nn.functional.pad

  Warning: this has not been thoroughly tested!
  &quot;&quot;&quot;
  if mode != 'constant':
    raise NotImplementedError(&quot;Only mode='constant' is implemented&quot;)
  assert len(pad) % 2 == 0
  assert len(pad) // 2 &lt;= input.ndim
  pad = list(zip(*[iter(pad)]*2))
  pad += [(0, 0)] * (input.ndim - len(pad))
  return lax.pad(
      input,
      padding_config=[(i, j, 0) for i, j in pad[::-1]],
      padding_value=jnp.array(value, input.dtype))

x = jnp.ones((2, 3))
y = jax_pad(x, (1, 1))
print(y)
# [[0. 1. 1. 1. 0.]
#  [0. 1. 1. 1. 0.]]

x = jax_pad(y, (-1, -1))
print(x)
# [[1. 1. 1.]
#  [1. 1. 1.]]

