vmap over axis with different sizes

[AlejandroTL]

Hello!

I have a function that receives as input a tuple of two jnp.ndarrays of different shapes, let's say (100x25) and (20x50). I want to apply vmap to iterate over the first dimension of the second element of my input, that is: (20x50). I want to iterate 20 times.

jax.vmap(my_function, in_axes=0)(tuple)

But I receive the following error that I don't know how to solve.

jax._src.traceback_util.UnfilteredStackTrace: ValueError: vmap got inconsistent sizes for array axes to be mapped: the tree of axis sizes is: ((100, 20),)

How can I specify that I want to iterate over the dimension of size 20 and not the one of size 100? I cannot do padding to make both inputs of same size.
If I do: jax.vmap(my_function, in_axes=(None, 0)(tuple), the error that I receive is different due to the input to the function:

vmap in_axes specification must be a tree prefix of the corresponding value, got specification (None, 0) for value tree PyTreeDef(((*, *),)).

Thanks!

[soraros]

Try the following as my_function is unary:

vmap(my_function, in_axes=((None, 0),))(tup)

[jakevdp]

The in_axes argument takes a list of PyTree objects (list, tuple, dict, etc.) whose structure corresponds to the structure of the associated inputs. So in your case you probably want in_axes=[(0, None)]. Here is a short example:

import jax
import jax.numpy as jnp

def my_function(params):
  x, y = params
  return x.sum() * y.sum()

params = (
    jnp.ones((100, 25)),
    jnp.ones((20, 50))
)

out = jax.vmap(my_function, in_axes=[(0, None)])(params)
print(out.shape)
# (100,)

[AlejandroTL]

Thank you both of you!

For the function jax.grad I have exactly same problem. My inputs are two tuples and I just want to take gradient with respect to one of the elements of the second tuple. I have tried using argnums=[(None,0)] but it doesn't work. Do you know how can this be solved?

[soraros]

@AlejandroTL Could you maybe give a more concrete example of what you want to achieve?

[NeilGirdhar]

@AlejandroTL I don't think grad supports that. I think you'll need to wrap your function in another one where what you want the gradient of is isolated. Maybe ask a separate question?

[AlejandroTL]

Hello!

Thank you all for your answers @soraros @NeilGirdhar. I have several problems using as input a tuple, so I decided to use two separated inputs x,y. The function that I'm using inside the vmap is a lambda function that uses jax.grad:

jax.vmap(lambda x, y: jax.grad(self.function, argnums=1)(params, x, y), in_axes=[None, 0])(x, y)

What I want is to take the gradients with respect to x, that's why argnums=1 and doing it iterating over the 0 dimension of y, that's why in_axis=[None,0].

The problem is that at the end I'm not obtaining a jnp.ndarray of dimensions y.shape[0] x some_dimension, which makes me think that I'm not really iterating over that dimension with the vmap instruction. Instead, my result is of dimensions x.shape[0] x some_dimension. Is there anyway in which I can take gradients with respect to x and iterate over y, so that I obtain at the end a stacked vector of all gradients w.r.t. x for each y?

I don't know if I explain myself properly.

Thanks!