darshandeshpande / jax-models

Unofficial JAX implementations of deep learning research papers

Hi 👋, I'm Darshan!

ML Researcher and a Mentor at TFUG Mumbai

• 🔭 I’m currently working on explainability in dialogue systems and causal language models.

• 🌱 I'm interested in argumentation, interpretability, and commonsense reasoning in large language models

• 🔎 I actively do research in the NLP subspace. Find my research on Google Scholar

• 📝 I write reports and articles on Weights & Biases | Medium

• 📫 Reach me at: [email protected]

• ⚡ Fun fact: I am a multicopter enthusiast

Tools and Libraries

jax-models's Issues

Missing Axis Swap in ExtractPatches and MergePatches

In patch_utils.py, the modules ExtractPatches and MergePatches are missing an axis swap between the reshapes, resulting in the extracted patches becoming horizontal stripes. For example, if we follow the code in ExtractPatches:

>>> 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))

For MergePatches, this should be:

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))