>>> inputs = jnp.arange(16).reshape(1, 4, 4, 1)
>>> inputs[0, :, :, 0]
DeviceArray([[ 0, 1, 2, 3],
[ 4, 5, 6, 7],
[ 8, 9, 10, 11],
[12, 13, 14, 15]], dtype=int32)
>>> patch_size = 2
>>> batch, height, width, channels = inputs.shape
>>> height, width = height // patch_size, width // patch_size
>>> x = jnp.reshape(inputs, (batch, height, patch_size, width, patch_size, channels))
>>> x = jnp.reshape(x, (batch, height * width, patch_size ** 2 * channels))
>>> x[0, 0, :]
DeviceArray([0, 1, 2, 3], dtype=int32)
We see that the first patch extracted is not the patch containing [0, 1, 4, 5], but the horizontal stripe [0, 1, 2, 3].
To fix this problem, we should add an axis swap. For ExtractPatches, this should be:
batch, height, width, channels = inputs.shape
height, width = height // patch_size, width // patch_size
x = jnp.reshape(
inputs, (batch, height, patch_size, width, patch_size, channels)
)
x = jnp.swapaxes(x, 2, 3)
x = jnp.reshape(x, (batch, height * width, patch_size ** 2 * channels))
batch, length, _ = inputs.shape
height = width = int(length**0.5)
x = jnp.reshape(inputs, (batch, height, width, patch_size, patch_size, -1))
x = jnp.swapaxes(x, 2, 3)
x = jnp.reshape(x, (batch, height * patch_size, width * patch_size, -1))