Jax vmap for simple array update

user_16004356 0
Viewed 248

I am new to Jax, and I am working on transforming someone else's code, which used the numba "fastmath" feature and relied on many nested for-loops without much performance loss. I am trying to recreate the same behavior using Jax's vmap function. However, I am currently struggling a lot with some fundamental questions. Here's a dumbed-down example of what I am trying to vectorized using vmap:

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

a = jnp.arange(20).reshape((4, 5))
b = jnp.arange(5)
c = jnp.arange(4)
d = jnp.zeros(20)
e = jnp.zeros((4, 5))

for i in range(a.shape[0]):
    for j in range(a.shape[1]):
        a = jax.ops.index_add(a, jax.ops.index[i, j], b[j] + c[i])
        d = jax.ops.index_update(d, jax.ops.index[i*a.shape[1] + j], b[j] * c[i])
        e = jax.ops.index_update(e, jax.ops.index[i, j], 2*b[j])

How would I rewrite such a code using vmap? While this code would be relatively easy to vectorize manually, I wish to understand better how vmap works and hope that any answer would help me. The docs don't seem to really help me right now. I really appreciate any help you can provide.

jax python
1 Answers

Here is how you might achieve roughly the same computation using vmap:

from jax import vmap, partial

@partial(vmap, in_axes=(0, None, 0))
@partial(vmap, in_axes=(0, 0, None))
def f(a, b, c):
  return a + b + c, b * c, 2 * b

a, d, e = f(a, b, c)
d = d.ravel()
user_29378310
Related