jax.vmap output: an explanation is required to get a 2D output array.

[jecampagne]

Hello,
Here is a snippet that exhibit different vmap outputs, with one that I do not expect and I do not manage to get it right neither. So I need help. Thanks

import jax.numpy as jnp

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

def fext1(kpos):
  return jnp.asarray(-1.)

def fext2(kpos):
  return -jnp.ones_like(kpos.x)

def fext3(kpos):
  return kpos.x + kpos.y  

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

Then with

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

I get

DeviceArray([-1., -1., -1., -1., -1., -1., -1., -1., -1., -1.],            dtype=float64, weak_type=True)

which I do not expect, while with

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

I get what I was expecting as a 10x10 array

DeviceArray([[-1., -1., -1., -1., -1., -1., -1., -1., -1., -1.],
             [-1., -1., -1., -1., -1., -1., -1., -1., -1., -1.],
             [-1., -1., -1., -1., -1., -1., -1., -1., -1., -1.],
             [-1., -1., -1., -1., -1., -1., -1., -1., -1., -1.],
             [-1., -1., -1., -1., -1., -1., -1., -1., -1., -1.],
             [-1., -1., -1., -1., -1., -1., -1., -1., -1., -1.],
             [-1., -1., -1., -1., -1., -1., -1., -1., -1., -1.],
             [-1., -1., -1., -1., -1., -1., -1., -1., -1., -1.],
             [-1., -1., -1., -1., -1., -1., -1., -1., -1., -1.],
             [-1., -1., -1., -1., -1., -1., -1., -1., -1., -1.]],            dtype=float64)

and the fext3 is ok too.

So, it seems after some other trials that the output of the function fext(kpos) should be of a certain type, but I do not manage to get it the correct one. Any idea is welcome.

[soraros]

Roughly, shape signature of the functions are:

# postulate
#   X: tuple[int, ...]
#   n: int
#   x, y: ndarray
#   x.shape == y.shape

# X, X -> ()
# (n, *X), (n, *X) -> (n,) under vmap
def fext1(x, y):
  return jnp.asarray(-1.)

# X, X -> X
# (n, *X), (n, *X) -> (n, *X) under vmap
def fext2(x, y):
  return -jnp.ones_like(x)

# X, X -> X
# (n, *X), (n, *X) -> (n, *X) under vmap
def fext3(kpos):
  return x + y

[jecampagne]

So, if I want to make f(x,y) returning a number which is the result of computations with kpos.x and kpos.y, I should multiply it by jnp.ones_like(kpos.x) to get it right?

[soraros]

Only if the shape dependency is not clear (or clear but the default doesn't suit your need) as in fext1.

[jecampagne]

1. So, is there a way to force the signature "(n, *X)" you mentionned?

2. Bye the way how do u have guessed the output signature, is there a JAX tools?

Thanks

[soraros]

There are a few simple shape inference rules:

• Input shape is determined by input argument
• For each primitive function (think lax functions), the value map f = lambda a: ...; b also defines (corresponding to) a shape map f: a.shape -> b.shape. Implemented as abstract_eval rules in JAX
• For functions f: b -> c and g: a -> b, the composition f . g has shape a -> c
• vmap (roughly) maps a function f: X, ... -> R into a function vmap(f): (n, *X), ... -> (n, *R)

Now, take fext1 as an example, we can get its shape signature as follows:

1. Input shape is the shape of whatever array you pass to it, so x, y: X, and fext1: X, X -> ?
2. Shape signature of asarray is (roughly) N -> N; its input -1. has shape (); the output asarray(-1.) has shape ()
3. The function fext1 has shape X, X -> ()

Applying the vmap rule gives you (n, *X), (n, *X) -> (n,). There is no guess work involved.

Could you elaborate on what do you mean by enforcing shape signature (n, *X)?

[soraros]

You could also make use of jax.eval_shape.

[jecampagne]

Thanks,
See discussion #14057 may be the example above will becomes clearer (hope so...)