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):
"""JAX implementation of torch.nn.functional.pad
Warning: this has not been thoroughly tested!
"""
if mode != 'constant':
raise NotImplementedError("Only mode='constant' is implemented")
assert len(pad) % 2 == 0
assert len(pad) // 2 <= 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.]]