How to vmap? - additional and different dimensional arrays

[harsh306]

I have two sets of parameters for which I want to compute gradients differently, and they could be of different size, so was wondering what would be the best way to do it?

import jax.numpy as jnp
from jax import grad, jit, vmap
import jax

x = jnp.array([0.5])
y = jnp.array([0.5 , 0.2])

def func(x, y, inputs, outputs):
    return inputs*jnp.square(x) +  outputs*jnp.square(y)

obj_grad = vmap(grad(func, argnums=[0]), in_axes=(0, None, None, None)) 
obj_grad(x, y, 3.0, 1.0)

Error:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-438-a925f04aa3d6> in <module>
----> 1 obj_grad(x, y, 3.0, 1.0 )

~/anaconda3/envs/cvx/lib/python3.7/site-packages/jax/api.py in batched_fun(*args)
    885     flat_fun, out_tree = flatten_fun_nokwargs(f, in_tree)
    886     in_axes_flat = flatten_axes("vmap in_axes", in_tree, in_axes)
--> 887     _ = _mapped_axis_size(in_tree, args_flat, in_axes_flat, "vmap")
    888     out_flat = batching.batch(flat_fun, args_flat, in_axes_flat,
    889                               lambda: flatten_axes("vmap out_axes", out_tree(),

~/anaconda3/envs/cvx/lib/python3.7/site-packages/jax/api.py in _mapped_axis_size(tree, vals, dims, name)
    929                   size)
    930                 for size, idxs in sizes.items()]
--> 931       raise ValueError(msg.format("\n".join(lines1 + ["so"] + lines2))) from None
    932     else:
    933       sizes = [x.shape[d] if d is not None else None for x, d in zip(vals, dims)]

ValueError: vmap got inconsistent sizes for array axes to be mapped:
arg 0 has shape (1,) and axis 0 is to be mapped
arg 1 has shape (2,) and axis 0 is to be mapped
arg 2 has shape () and axis None is to be mapped
arg 3 has shape () and axis None is to be mapped
so
arg 0 has an axis to be mapped of size 1
arg 1 has an axis to be mapped of size 2

[harsh306]

I think I misunderstood the use of VMAP. Few things I learned:

• Output of the function to apply grad should output a scalar value. Not a vector, even of shape (1, ).
• You can then pass all your parameters as a list to the objective function of your choice, (yes of arbitrary shapes to compute grad).
• Even you can provide multiple lists of parameters or jax arrays. (For example following works)
• Advantage of sending multiple lists is you can choose which args to compute gradients on.

import jax.numpy as jnp
from jax import grad, jit, vmap
from jax import random
import jax
import functools
p = jnp.array([1.0, 2.0, 3.0])
b = jnp.array([1.0])
params = [p, b]

def func(p, b, inputs, outputs):
    return jnp.mean(inputs*jnp.square(p) + outputs*jnp.square(b))

print(func(p, b, 3, 2)) # check this is scalar
grad_fun = grad(objective, argnums=[0, 1])
grad_fun(p, b, 3, 2)