double jax.vmap working as two nested loops, why single jax.vmap also working sometimes.

[jecampagne]

Well, here is a snippet to discuss a little on vmap (a complement of #14054)

import jax.numpy as jnp

class Position():
  def __init__(self, x,y):
    self.x =  jnp.asarray(x)
    self.y =  jnp.asarray(y)

class Obj():
  def __init__(self, sigma, flux):
    self.flux = flux
    self.sigma = sigma
    self._sigsq = self.sigma**2

  def _kValue(self, kpos):
        ksq = (kpos.x**2 + kpos.y**2) * self._sigsq
        return self.flux * jnp.exp(-0.5 * ksq)

make a list with two Objects

a_list = [Obj(sigma=1., flux=1.), Obj(sigma=2.,flux=0.1)]

Then, define a class that manipulate list of Objects

class A():
  def __init__(self,a):
    self.a_list = a

  def f(self, kpos):
      k_list = [obj._kValue(kpos) for obj in self.a_list]
      return jnp.prod(jnp.array(k_list))    #*jnp.ones_like(kpos.x) (see another post)

  def g(self, kpos):
      return self.a_list[0]._kValue(kpos)

Define also an array of positions

coords = jax.random.normal(jax.random.PRNGKey(10), shape=(5,5,2))*0.2

Let us predict what would do a double for loop mimicing the "A:f" function

for i in range(coords.shape[0]):
  for j in range(coords.shape[1]):
    print(i,j,jnp.prod(jnp.array([obj._kValue(Position(coords[i,j,0],coords[i,j,1])) for obj in a_list])))

One gets

0 0 0.0872824173694711
0 1 0.0833887684154781
0 2 0.06885071059512518
0 3 0.09556858419144411
0 4 0.07422189674104165
1 0 0.0687560349953797
1 1 0.08822973615774296
1 2 0.07913844117885234
1 3 0.08147450699633364
1 4 0.08645643145102813
2 0 0.08980484577095833
2 1 0.09236186377307067
2 2 0.06824952747512257
2 3 0.09415511467179831
2 4 0.0857002252450574
3 0 0.09197638674363777
3 1 0.05929086418838377
3 2 0.08768241912525869
3 3 0.08515955003365248
3 4 0.09517315229381747
4 0 0.0858689831775795
4 1 0.0799892467771078
4 2 0.0639629169171923
4 3 0.07358062559644361
4 4 0.08659691260134722

Of course this double for loop looks horrible for JAX, so let us try with vmap using the A:f fucntion

# instance of the A class with our list of Objects
tmpA = A(a_list)
#define a nested double vmap
vf = jax.vmap(jax.vmap(lambda i,j: tmpA.f(Position(coords[i,j,0],coords[i,j,1])), in_axes=(None, 0)), in_axes=(0, None))
# then apply it as
vf(jnp.arange(5), jnp.arange(5))

then one gets

DeviceArray([[0.08728242, 0.08338877, 0.06885071, 0.09556858, 0.0742219 ],
             [0.06875603, 0.08822974, 0.07913844, 0.08147451, 0.08645643],
             [0.08980485, 0.09236186, 0.06824953, 0.09415511, 0.08570023],
             [0.09197639, 0.05929086, 0.08768242, 0.08515955, 0.09517315],
             [0.08586898, 0.07998925, 0.06396292, 0.07358063, 0.08659691]],  dtype=float64)

which looks pretty good.

The point is that if one does

jax.vmap(lambda *args: tmpA.f(Position(*args)))(
        coords[..., 0], coords[..., 1]
    )

the result is completely wrong

DeviceArray([3.55459359e-06, 3.38167871e-06, 4.56791473e-06,
             3.87547072e-06, 2.79938140e-06], dtype=float64)

But comes (may be) the interesting part: If you look at the A:g function you will see that it uses only the first Object of the a_list. Now

vg = jax.vmap(jax.vmap(lambda i,j: tmpA.g(Position(coords[i,j,0],coords[i,j,1])), in_axes=(None, 0)), in_axes=(0, None))
vg(jnp.arange(5), jnp.arange(5))

gives

DeviceArray([[0.97316247, 0.96432075, 0.92807203, 0.99097576, 0.9421204 ],
             [0.92781666, 0.9752658 , 0.95428368, 0.95985216, 0.97131358],
             [0.97872336, 0.98423439, 0.92644561, 0.98802693, 0.96960845],
             [0.98341147, 0.90073607, 0.97405281, 0.96838192, 0.99015433],
             [0.96999002, 0.95632679, 0.91450409, 0.94048678, 0.97162903]],

AND INTERESTINGLY

jax.vmap(lambda *args: tmpA.g(Position(*args)))(
        coords[..., 0], coords[..., 1]
    )

gives the same correct result.

So, what is puzzling me is the difference of behavior between jax.vmap(lambda *args: tmpA.g... and jax.vmap(lambda *args: tmpA.f....

1. if one can gives a version with a single vmapinstead of the vf (vg) version that use 2 vmap, I would be glad, and
2. if one can gives a why the single vmap is correct with a unique Object in the list (contrary to the a_list containing 2 objects) then I would be glad too.

[soraros]

Ah, I see the problem.

1. outer-like

# postulate
a.shape == (m,)
b.shape == (n,)

vmap(vmap(f, in_axes=(None, 0)), in_axes=(0, None))(a, b)

which gives

[ f(a0, b0), f(a0, b1), ..., f(a0, bn);
  f(a1, b1), f(a1, b1), ..., f(a1, bn);
  ...
  f(am, bn), f(am, b1), ..., f(am, bn);
]

2. vectorize-like

# postulate
a.shape == b.shape == (m, n)

vmap(vmap(f, in_axes=(0, 0)), in_axes=(0, 0))(a, b)

which gives

[ f(a00, b00), f(a01, b01), ..., f(a0n, b0n);
  f(a10, b10), f(a11, b11), ..., f(a1n, b1n);
  ...
  f(am0, bm0), f(am1, bm1), ..., f(amn, bmn);
]

Now go back to your code

X, Y = coords[..., 0], coords[..., 1]

# scenario 1, correct
f1 = vmap(vmap(lambda i, j: f(X[i, j], Y[i, j]),
               in_axes=(None, 0)), in_axes=(0, None))
f1(jnp.arange(5), jnp.arange(5))

# scenario 2, incorrect
f2 = vmap(vmap(lambda x, y: f(x, y),
               in_axes=(None, 0)), in_axes=(0, None))
f2(X, Y)

# scenario 3, correct
f3 = vmap(vmap(lambda x, y: f(x, y),
               in_axes=(0, 0)), in_axes=(0, 0))
f3(X, Y)

Your scenario 2 is essentially

[ f(X[0,:], Y[0,:]), f(X[0,:], Y[1,:]), ..., f(X[0,:], Y[5,:]);
  f(X[1,:], Y[1,:]), f(X[1,:], Y[1,:]), ..., f(X[1,:], Y[5,:]);
  ...
  f(X[5,:], Y[0,:]), f(X[5,:], Y[1,:]), ..., f(X[5,:], Y[5,:]);
]

It happens to run without raising shape error only because X and Y are square matrices.

[soraros]

@jecampagne

jax.vmap(jax.vmap(lambda *args: tmpA.f(Position(*args)
                  in_axes=(0, 0)), in_axes=(0, 0))  # this line is redundant, as in_axes defaults to 0
         )(coords[..., 0], coords[..., 1])

[jecampagne]

Thanks a lot for your clear explanation.

[soraros]

@jecampagne You can guard against this kind of error by using assertions. Something like

class A:
  def f(self, kpos):
    assert kpos.x.shape == kpos.y.shape == ()
    k_list = [obj._kValue(kpos) for obj in self.a_list]
    return jnp.prod(jnp.array(k_list))

[soraros]

@jecampagne There is an even better trick to use. You could make Position a NamedTuple and greatly simplify the call-site, because NamedTuple is a PyTree and they are vmap compatible.

from typing import List, NamedTuple

import jax
import jax.numpy as jnp
from jax import vmap


class Position(NamedTuple):
  x: jax.Array
  y: jax.Array

class Obj:
  def __init__(self, sigma: float, flux: float):
    self.flux = flux
    self.sigma = sigma
    self._sigsq = sigma**2

  def k_val(self, kpos: Position):
    ksq = (kpos.x**2 + kpos.y**2) * self._sigsq
    return self.flux * jnp.exp(-0.5 * ksq)

class A:
  def __init__(self, a: List[Obj]):
    self.a_list = a

  def f(self, kpos: Position):
    return jnp.array([obj.k_val(kpos) for obj in self.a_list]).prod()

coords = 0.2 * jax.random.normal(jax.random.PRNGKey(10), shape=(2, 5, 5))
a = A([Obj(sigma=1., flux=1.), Obj(sigma=2., flux=0.1)])

vmap(vmap(a.f))(Position(*coords))

[jecampagne]

@soraros , I was wandering if you understand why the "A:g" function works with a single vmap (this is the point n°2 in my post) ?