[Python] Looking for an efficent JAX function to reconstruct an image from patches

I have a set of images in (c, h, w) jax arrays. These arrays have been converted to (patch_index, patch_dim) arrays where patch_dim == c * h * w.

I am trying to reconstruct the original images from the patches. Here is vanilla python code that works:

kernel = jnp.ones((PATCH_DIM, IMG_CHANNELS, PATCH_HEIGHT, PATCH_WIDTH), dtype=jnp.float32)

def fwd(x):
xcv = lax.conv_general_dilated_patches(x, (PATCH_HEIGHT, PATCH_WIDTH), (PATCH_HEIGHT, PATCH_WIDTH), padding='VALID')

# return channels last
return jnp.transpose(xcv, [0,2,3,1])

patches = fwd(bfrc)

patch_reshaped_pn_c_h_w = patch_reshaped_ph_pw_c_h_w = jnp.reshape(patches, (V_PATCHES, H_PATCHES, IMG_CHANNELS, PATCH_HEIGHT, PATCH_WIDTH))

# V_PATCHES == IMG_HEIGHT // PATCH_HEIGHT
# H_PATCHES == IMG_WIDTH // PATCH_WIDTH

reconstructed = np.zeros(EXPECTED_IMG_SHAPE)

for vpatch in range(0, patch_reshaped_ph_pw_c_h_w.shape[0]):
for hpatch in range(0, patch_reshaped_ph_pw_c_h_w.shape[1]):
for ch in range(0, patch_reshaped_ph_pw_c_h_w.shape[2]):
for prow in range(0, patch_reshaped_ph_pw_c_h_w.shape[3]):
for pcol in range(0, patch_reshaped_ph_pw_c_h_w.shape[4]):
row = vpatch * PATCH_HEIGHT + prow
col = hpatch * PATCH_WIDTH + pcol
reconstructed[0, ch, row , col] = patch_reshaped_ph_pw_c_h_w[vpatch, hpatch, ch, prow, pcol]

# This assert passes
assert jnp.max(jnp.abs(reconstructed - bfrc[0])) == 0


Of course this vanilla python code is very inefficient, but I am not sure how to convert the for loops into efficient JAX code. If anyone has suggestions, that would be greatly appreciated.

Continue reading...