How to vmap over this function?

[bigyankarki]

I have this function which works for single vector:

def board_mapping(dim):
    mapping = dim * dim * [0]
    turn = 1
    base_ind = 0
    for i in range(dim):
        ind = base_ind
        for j in range(dim - 1):
            mapping[ind] = i + j * dim

            if turn + j >= dim:
                offset = dim - (turn + j - dim + 1)
            else:
                offset = turn + j

            ind += offset
        mapping[ind] = i + (j + 1) * dim

        turn += 1
        base_ind += turn

    return np.array(mapping)

mapping = board_mapping(dim=4)
reverse_mapping = ((4 ** 2) - 1) - mapping

def vec_to_board(vector, player, dim, reverse=False):
    player_board = np.zeros(dim * dim)
    player_pos = np.argwhere(vector == player)
    if not reverse:
        player_board[mapping[player_pos.T]] = 1
    else:
        player_board[reverse_mapping[player_pos.T]] = 1
    return np.reshape(player_board, [dim, dim])

Sample input:

vec = [1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 2, 2]
p1 = vec_to_board(vec, 1, 4)
p2 = vec_to_board(vec, 2, 4)

Output

p1 = [[1. 1. 0. 0.]
 [1. 0. 0. 0.]
 [0. 0. 0. 0.]
 [0. 0. 0. 0.]]

p2 = [[0. 0. 0. 0.]
 [0. 0. 0. 0.]
 [0. 0. 0. 1.]
 [0. 0. 1. 1.]]

However, I want it to work for a batch of vectors.

What I have tried so far:

states = jnp.array([[1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 2, 2], [1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 2, 2]])
 batch_size = 1
b_states = vmap(vec_to_board)((states, 1, 4), batch_size)

This doesn't work. However, if I understand correctly vmap should be able to handle this transformation for batches?

[jakevdp]

Can you update your question with the complete code needed to run your function, including example inputs? In particular, mapping and reverse_mapping are undefined, so it's difficult to guess what your function does or what error you're seeing when you attempt to vmap. Thanks!

[bigyankarki]

I have updated the code above to reflect the changes. Heres a working code: https://colab.research.google.com/drive/1Jit85WeDRLTgUJic9rE6NE04x9Qa29oa?usp=sharing

[jakevdp]

Thanks. Unfortunately, your function is using a number of constructs that are incompatible with JAX transforms because they rely on dynamically-shaped arrays (for example, the output of argwhere is dynamic). You won't be able to use vmap on this function unless you come up with an alternative approach that's compatible with dynamic shapes.

[jakevdp]

See also https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#in-place-updates for information on doing in-place updates in JAX, which is another incompatible construct your function uses.