Updating slices in a pytree with slices from another pytree of the same shape

[StoneT2000]

I currently have a PyTree with the following shape

JuxAction(
    factory_action=ShapedArray(int8[10,2,6]),
    unit_action_queue=UnitAction(
        action_type=ShapedArray(int8[10,2,200,20]),
        direction=ShapedArray(int8[10,2,200,20]),
        resource_type=ShapedArray(int8[10,2,200,20]),
        amount=ShapedArray(int16[10,2,200,20]),
        repeat=ShapedArray(int16[10,2,200,20]),
        n=ShapedArray(int16[10,2,200,20])
    ),
    unit_action_queue_count=ShapedArray(int8[10,2,200]),
    unit_action_queue_update=ShapedArray(bool[10,2,200])
)

Notably each leaf has a shared dimensions of (10, 2) = (batch_size, team_id) (this is for a game environment with two players).

Suppose I have two of these, jux_action_0, jux_action_1, however the following doesn't work

jux_action: JuxAction = jax.tree_map(lambda x, y : x.at[:, 1].set(y.at[:, 1]), jux_action_0, jux_action_1)

I was hoping this would auto set the 2nd element in the 2nd dimension with what is stored in jux_action_1 and leave the rest as whatever was in jux_action_0. However I get the following error

TypeError: Value '_IndexUpdateRef(Array([[[ 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, -1, -1],
        [-1, -1, -1, -1, -1, -1]],

       [[ 1, -1, -1, -1, -1, -1],
        [-1, -1, -1, -1, -1, -1]]], dtype=int8), (slice(None, None, None), 0))' with dtype object is not a valid JAX array type. Only arrays of numeric types are supported by JAX.

This seems to be because the return value of the lambda passed to tree_map is the output of .set. Is there another way to set the value of a slice with another slice and return it in tree_map?

[jakevdp]

I think you want this:

jux_action: JuxAction = jax.tree_map(lambda x, y : x.at[:, 1].set(y[:, 1]), jux_action_0, jux_action_1)

or if you want to use y.at, you can do this:

jux_action: JuxAction = jax.tree_map(lambda x, y : x.at[:, 1].set(y.at[:, 1].get()), jux_action_0, jux_action_1)

[StoneT2000]

oh my that's such a silly mistake 😅 thanks!