Applying a scalar-input root-finding function to elements of an array - using vmap passes the entire array, rather than each scalar value.

[jackbashford]

I'm currently doing some work with this JAX snippet, which is intended to map a 1D array element-by-element, and depending on whether the element satisfies a certain criterion function, it will run a root-finding function:

rootFinder = lambda M: brentq(f, a=constant, b=M) # this line is simplified

mapFn = lambda M: jax.lax.cond(f(M), lambda: M, lambda: rootFinder(M))

return jax.vmap(mapFn)(arr)

Both f and rootFinder require their input values to be of a scalar type - that is, they can't take the arrays as input (which, when I put print statements in, is what appears to be happening while JAX is tracing the vmap). Is it possible to vectorise this transformation using vmap as I'm trying to do? Or do I need to use a standard for loop or other tools to try and implement this in an efficient manner (without lots of .at.sets everywhere?). This function will be run many thousands of times inside an iterative solution, which is why I'm trying to get efficiency out of it (not exactly sure whether the overhead of using brentq is worth it compared to a less efficient root-finding algorithm from jaxopt, any advice on that front would also be greatly appreciated).

[anh-tong]

Hi,

According to JAX docs, when using vmap over the predicate f, jax.lax.cond will turn into jax.lax.select which evaluates both branches. Therefore, rootFinder will run over the whole array. If it is costly to run, this may be not the best practice.

As you mentioned, using loop can be a solution. Here is a quick comparison between using vmap and jax.lax.scan.

import jax
import jax.numpy as jnp
import jax.random as jrandom

# make a somewhat expensive function
def dummy_fn(m):
  return jnp.sum(jnp.ones((1000, )))

# dummy predicate
def f(m):
  return m > 0

map_fn = lambda m: jax.lax.cond(f(m), lambda : m, lambda : dummy_fn(m))

Let's say the 1D array is generated as

# since this is randomly generated, there is about 50% the the predicate is True
m = jrandom.normal(key=jrandom.PRNGKey(0), shape=(1000,))

using vmap

%%timeit
jax.vmap(map_fn)(m)
# 6.62 ms ± 1.26 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)

Now, using jax.lax.scan

# scan function for jax.lax.scan
def body_fn(current, x):
  ret = map_fn(x)
  return ret, ret

%%timeit
jax.lax.scan(body_fn, m[0], m)
# 740 µs ± 252 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

You can check out a google colab here.

In this case, jax.lax.scan is the winer. However, the situation depends on the overhead of dummy_fn and how frequent the predicate selects the expensive branch.

[jackbashford]

Hmm, I might not have the root cause of the issue then. When using vmap I will run into the issue of the root finding function being run when it's not needed (as expected). But when I use scan, I either get an error that a float is being treated like an array somewhere (if I use pure_callback), or I get the issue that the function is being run when it's not needed (this function will unfortunately error when that happens) - and that's when I don't use a callback. So the 'best shot' code I currently have is:

c1_map = lambda M: jax.lax.cond(predicate(M), lambda: M, lambda: jax.pure_callback(rootFinder, M, M))

def body_fn(_, x):
    ret = c1_map(x)
    return ret, ret

return jax.lax.scan(body_fn, grid[0], grid)

This gives me back:

jaxlib.xla_extension.XlaRuntimeError: INTERNAL: Generated function failed: CpuCallback error: AttributeError: 'float' object has no attribute 'shape'

Which is almost certainly something to do with the root finding algorithm expecting an array as input (but wrapping the passed value in jnp.array does nothing) - unsure how to progress beyond this point, but thanks for the help regarding scan so far! If I can share anything that might make problem-solving this easier, I'd be happy to provide mini reproducible demos.

[jackbashford]

I managed to determine the cause of the issue - when JAX is checking the return type of a callback, it tries to compare shapes - but a float doesn't have a shape attribute. So the fix was to wrap jnp.array to my returned value inside a proxy call to rootFinder. After doing this, scan works just fine and I can even wrap the entire call in a while_loop and the performance is almost identical to the list comprehension solution I was using.

Thanks for the assistance regarding vmap vs scan @anh-tong !

[jakevdp]

Note that in vmap, the shape of a scalar will be (), although the __repr__ of a scalar will show the entire vmapped array. For example:

import jax

def f(x):
  print("shape =", x.shape)
  print(x)
  return x

x = jax.numpy.arange(10)

jax.vmap(f)(x)
# shape = ()
# Traced<ShapedArray(int32[])>with<BatchTrace(level=1/0)> with
#   val = Array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], dtype=int32)
#   batch_dim = 0

So I suspect in your case you actually were seeing a scalar value, but the repr was confusing that.

[jackbashford]

That's definitely something that's useful to know, thanks for the info - didn't realise the __repr__ behaved that way.