Vectorization and Named Arguments in JAX

[GoktugGuvercin]

Hello JAX Community,

While working on vectorization in JAX, I noticed something confusing and I am here to ask it: I was planing to vectorize $pix.rotate()$ function in Deepmind's image processing library PIX by using $jax.vmap()$. I wrote the vectorization code like this:

vmap_rotate = jax.vmap(pix.rotate, in_axes=(0, 0), out_axes=0)

My aim was to make $pix.rotate()$ take multiple images and multiple rotation radians as input, and apply the rotation of each image with respective radians. In other words, I was trying to make the function adapted to batching. It works quite well, it seems to have no mistake. However, after vectorization, I cannot use named argument mode of $pix.rotate()$ function. In other words, the attempt to call vmap_rotate(images, radians, mode="nearest") throws error. I am thinking that it is about the vectorization. After vectorizing a function, are not we allowed to use named arguments ?

[mattjj]

Hey Göktuğ!

Indeed, the function returned by vmap (here named vmap_rotate) has an often-surprising interaction with arguments passed as keywords. The reason is ultimately because in_axes refers to parameters by position, and we don't attempt to identify those positions with parameters' names (originally because there wasn't a robust mechanism to do that in Python, though it's gotten a bit better with more recent versions of Python, while still imperfect). So instead of arguments passed by keyword being identified with any of the in_axes, they're just always mapped over their leading axis. (I'd like to revise that behavior, but haven't gotten to it, in part because I know it'll break existing code and deprecation cycles are annoying.)

Luckily there's an easy workaround: just write a wrapper which ensures the vmapped function is always called with positional arguments:

_vmap_rotate = jax.vmap(pix.rotate, in_axes=(0, 0), out_axes=0)  # note leading underscore

def vmap_rotate(images, radians, mode):
  return _vmap_rotate(images, radians, mode)

Here the vmap_rotate function just serves to canonicalize how arguments are passed to _vmap_rotate, the vmapped function.

WDYT?

[GoktugGuvercin]

After your answer, I checked $jax.vmap()$ documentation and I noticed that as you said, in_axes is configured to perceive positional arguments to the function that tries to be v-mapped. Besides, your advice is so elegant. The wrapper can directly solve the problem. Thank you so much.