lax.conv and vmap

[kyunghyuncho]

lax.conv is already defined to operate on a batched set of images, as stated e.g. in https://jax.readthedocs.io/en/stable/notebooks/Common_Gotchas_in_JAX.html?highlight=lax.conv#lax.conv-and-lax.conv_with_general_padding. this makes it a bit awkward to use it together with vmap, as it becomes necessary to expand_dims any single example with an extra dimension at the beginning and squeeze the outcome to mimic the single-example computation using lax.conv and to vmap it. for instance, see

import jax

from jax import numpy as jnp
import numpy as onp

image = onp.random.randn(3,13,17)
batched_image = onp.random.randn(7,3,13,17)
weight = onp.random.randn(5,3,2,2)

def myconv(im, we):
    return jax.lax.conv(im, we, (1,1), 'SAME')

myconv_ = jax.jit(myconv)

# myconv_(image, weight).shape <= error TypeError: convolution requires lhs and rhs ndim to be equal, got 3 and 4.
x = myconv_(batched_image, weight)

myconv_vmap = jax.vmap(myconv, (0, None))

# myconv_vmap(image, weight).shape <= error TypeError: convolution requires lhs and rhs ndim to be equal, got 2 and 4.
y = myconv_vmap(onp.expand_dims(batched_image,1), weight) 

print('x.shape=', x.shape, 'y.shape=', y.shape, 'difference',((x - y.squeeze()) ** 2).sum())

perhaps i'm missing some thoughts behind this design decision behind lax.conv.

[jakevdp]

My understanding is that the API of jax.lax.conv directly mirrors that of XLA:Conv. Since XLA itself doesn't have any generalized vmap-like transformation functionality, XLA's API builds-in this batching behavior.

Have you seen jax.scipy.signal.convolve/jax.scipy.signal.convolve2d and numpy.convolve? These mirror the (non-batched) convolution semantics of numpy and scipy, and might be more amenable to what you have in mind.

[mattjj]

I think the numpy/scipy functions might not support channels (e.g. in how the singleton axes are added here).

Actually this is an old wart in JAX, and one we've never gotten around to fixing: see #381 (and probably several duplicates)! The right thing to do is adjust our conv primitive not to require a batch dimension. We've never done it, to our great shame, and that is an impediment to the vmap dream of liberating ourselves from batch dimensions. (Tangentially, another impediment was batch norm, but vmap recently gained support for naming the mapped axis and applying collectives, so that's no longer an issue!)

We should finally fix this conv batch issue once and for all. I'll mention it at our team meeting today and see if we can get any takers. Let's track in #381.

[georgedahl]

Cast off the legacy of shame!

[cwhy]

May I know why #381 is closed now? I don't think this issue is related to stax... I am just trying to write composed functions with jax that supports vmap, which is blocked by convolution here.

[cwhy]

Would the work-around by OP have any potential performance penalty, if not, would that be a good addition to jax.image?🤔

[kyunghyuncho]

thanks, @jakevdp and @mattjj! it'll be great to have it fixed to realized the "dream of liberating ourselves from batch dimensions", which is a great dream.