Trying to implement some tf functions into jax

[bigyankarki]

I have this function in tensorflow:

def tf_vec_to_board(vector, player, batch_size):
    board = tf.tile(tf.expand_dims(tf.constant(range(cw.BOARD_SIZE**2),
        dtype=tf.int64), 0), [batch_size, 1])
    tiled_mapping = tf.tile(tf.expand_dims(tf_mapping, 0),
            [batch_size, 1])
    player_pos = tf.where(tf.constant(vector, dtype=tf.int64) == player,
            None, None)
    locations = tf.expand_dims(tf.reshape(tf.gather_nd(tiled_mapping,
        player_pos), [batch_size, -1]), axis=2)
    board = tf.cast(tf.reduce_any(tf.subtract(tf.expand_dims(board, axis=1),
        locations) == 0, axis=1), tf.int64)
    return tf.reshape(board, [batch_size, cw.BOARD_SIZE, cw.BOARD_SIZE])

And, I have refactored the code to Jax as following:

def unbatched_gather_nd(params, indices):
    return params[tuple(jnp.moveaxis(indices, -1, 0))]

def batched_gather_nd(params, indices):
    return vmap(unbatched_gather_nd, (0,0), 0)(params, indices)

def vec_to_board(vector, player, batch_size):
    print("player: ", player)
    print("vector: ", vector)
    print(vector.shape)
    print("batch_size: ", batch_size)

    board = jnp.tile(jnp.expand_dims(jnp.array(range(cw.BOARD_SIZE**2), dtype=np.int64), 0), [batch_size, 1])
    print("the board is :", board)

    tiled_mapping = jnp.tile(jnp.expand_dims(tf_mapping, 0), [batch_size, 1])
    print("Tiled Mapping: ", tiled_mapping)

    player_pos = jnp.where(jnp.array(vector, dtype=np.int64) == player, None, None)
    print("player pos: ", player_pos)

    locations = jnp.expand_dims(jnp.reshape(batched_gather_nd(tiled_mapping, player_pos), [batch_size, -1]), axis=2)
    print("locations: ", locations)

    board = jnp.array(jnp.any(jnp.subtract(jnp.expand_dims(board, axis=1), locations) == 0, axis=1)).astype(np.int64)
    print("final board: ", player)

    return jnp.reshape(board, [batch_size, cw.BOARD_SIZE, cw.BOARD_SIZE])

However, I keep on getting an error:

...
  File "/Users/bigyankarki/Desktop/bigyan/cc/chinese-checkers/python/nn_player.py", line 54, in expand_actions
    p1 = pw.vec_to_board(jnp.array(jnp.expand_dims(input_rep, axis=0)).astype(np.int64), 1, 1)
  File "/Users/bigyankarki/Desktop/bigyan/cc/chinese-checkers/python/pywrapper.py", line 189, in vec_to_board
    locations = jnp.expand_dims(jnp.reshape(batched_gather_nd(tiled_mapping, player_pos), [batch_size, -1]), axis=2)
  File "/Users/bigyankarki/Desktop/bigyan/cc/chinese-checkers/python/pywrapper.py", line 171, in batched_gather_nd
    return vmap(unbatched_gather_nd, (0,0), 0)(params, indices)
ValueError: vmap got inconsistent sizes for array axes to be mapped:
the tree of axis sizes is:
(1, (3, 3))

The inputs are:
player: 1
vector: [[1 1 1 0 0 0 0 0 0 0 0 0 0 2 2 2]]
batch_size: 1

Output so far:

Tiled mapping: [[ 0  4  1  8  5  2 12  9  6  3 13 10  7 14 11 15]]
player pos:  (DeviceArray([0, 0, 0], dtype=int32), DeviceArray([0, 1, 2], dtype=int32))```

Tensorflow output, and expected output:
```board:  tf.Tensor([[ 0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15]], shape=(1, 16), dtype=int64)
tiled mapping: tf.Tensor([[ 0  4  1  8  5  2 12  9  6  3 13 10  7 14 11 15]], shape=(1, 16), dtype=int64)
player_pos:  tf.Tensor(
[[ 0 11]
 [ 0 13]
 [ 0 15]], shape=(3, 2), dtype=int64)
locations:  tf.Tensor(
[[[10]
  [14]
  [15]]], shape=(1, 3, 1), dtype=int64)
final board:  tf.Tensor([[0 0 0 0 0 0 0 0 0 0 1 0 0 0 1 1]], shape=(1, 16), dtype=int64)```

AFAIK, theres problem in batched_gather_nd function, but I cannot understand the error that I am getting.

[YouJiacheng]

Try:

def batched_gather_nd(params, indices):
    return vmap(unbatched_gather_nd, (None, 0), 0)(params, indices)

[bigyankarki]

This seems to work for me:

board = jnp.tile(jnp.expand_dims(jnp.array(range(cw.BOARD_SIZE**2), dtype=np.int64), 0), [batch_size, 1])
    tiled_mapping = jnp.tile(jnp.expand_dims(tf_mapping, 0), [batch_size, 1])
    player_pos = jnp.where(jnp.array(vector, dtype=np.int64) == player, None, None)
    player_pos_stacked = jnp.stack(player_pos, axis=-1)
    locations = jnp.expand_dims(jnp.reshape(unbatched_gather_nd(tiled_mapping, player_pos_stacked), [batch_size, -1]), axis=2)
    board = jnp.array(jnp.any(jnp.subtract(jnp.expand_dims(board, axis=1), locations) == 0, axis=1)).astype(np.int64)
    return jnp.reshape(board, [batch_size, cw.BOARD_SIZE, cw.BOARD_SIZE])