Padding doesn't work in jitted mode

Hi,

Thanks for the work on the library. For a project I’m needing to implement a padding on the last dimension of a tensor. Everything works fine when the function is not jitted, but I get a conversion error when I jit it. It seems to me like it is coming from a bug in the lax_numpy internals (this line), but maybe I’m missing something…

Adrien

The code is as follows:

from jax import jit as jjit
import jax.numpy as jnp
import numpy as np

vals = np.random.randn(50, 100)

def pad_last_dim(array, pad_size):
    ndim = jnp.ndim(array)
    npad = jnp.zeros((ndim, 2), dtype=jnp.int32)
    axis = ndim - 1
    npad = ops.index_update(npad, ops.index[axis, 1], pad_size)
    return jnp.pad(array, npad, 'constant', constant_values=0)

print(pad_last_dim(vals, 1))  # All good
print(jjit(pad_last_dim)(vals, 1))  # raises

[AdrienCorenflos]

Thanks a lot for that Matt. Bit late for me but I’ll come back to it tomorrow. I think you should forbid lists and ndarrays tbh. Would be consistent with tf where they consider lists to be a collection of some kind.

In the meantime I like your solution and I’ll use it (I don’t need the ndarray, just was easier to set the padding)

[AdrienCorenflos]

from jax import jit as jjit
import jax.numpy as jnp
from jax import ops
import numpy as np

vals = np.random.randn(50, 100)

def pad_last_dim(array, pad_size):
    ndim = jnp.ndim(array)
    npad = [[0, pad_size]] * ndim
    return jnp.pad(array, npad, 'constant', constant_values=0)

print(pad_last_dim(vals, 1))  # All good
print(jjit(pad_last_dim, static_argnums=(1,))(vals, 1))

Your solution worked for me, thanks a lot!