Using vmap() with batched lists of GraphTuples

[miamirabelli]

Hi all,

To preface: I'm very new to JAX, so I apologize if this is a silly question. I've looked around for help re: GraphTuples and vmap (especially #16641 ), but I haven't had much luck. I'd appreciate any help!

I'm writing a two-layer GNN that currently predicts pretty well. However, it runs somewhat slow, so I'm looking for ways to optimize it on my GPU. I decided to implement batching during training; previously, my code just trained one time-series forecast window of data GraphTuples at a time. Now, I want to group batch_size time-series forecast windows into some number of batches, and use jax.vmap() to train the windows in parallel.

Here's a snippet of my code. I pass in my input and target batches, then use jax.vmap() to batch them into their time-series forecast windows:

def loss_fn(params, input_batch_graphs, target_batch_graphs):
        curr_state = state.replace(params=params) 
        # compute loss by window
        def compute_loss(input_window_graphs, target_window_graphs):
            # Compute loss for n rollout steps.
            x1_loss, x2_loss, _ = rollout_loss(
                state=curr_state, input_window_graphs=input_window_graphs, 
                target_window_graphs=target_window_graphs, n_rollout_steps=n_rollout_steps, 
                rngs=rngs)
            total_loss = x1_loss + x2_loss
            return total_loss, x1_loss, x2_loss
        
        # defines in_axes as a tuple of GraphTuples, specifying which axes to vectorize over
        in_axes = (
            jraph.GraphsTuple(nodes=1, edges=1, receivers=1, senders=1, globals=1, n_node=None, n_edge=None),
            jraph.GraphsTuple(nodes=1, edges=1, receivers=1, senders=1, globals=1, n_node=None, n_edge=None))
        
        # applies the loss fn to every window in the batches
        batch_losses, batch_x1_losses, batch_x2_losses, batch_pred_nodes = jax.vmap(
            lambda x, y: compute_loss(x, y), in_axes=in_axes)(input_batch_graphs, target_batch_graphs)

Here's a very small version of what input_batch_graphs is structured like. In here, batch_size = 2, and each graph contains 6 nodes, each with 2 layers:

[GraphsTuple(nodes=Array([[ 1.1636245 ,  1.0790482 ],
       [-1.293495  , -0.35579923],
       [ 0.37819758,  0.18474741],
       [ 0.37264496, -1.640342  ],
       [ 0.20188326, -1.047298  ],
       [ 1.1561959 ,  3.074735  ]], dtype=float32), 
 edges=Array([[ 0.], [ 1.], [ 2.], [-1.], [-2.],
       [ 0.], [ 1.], [ 2.], [-1.], [-2.],
       [ 0.], [ 1.], [ 2.], [-1.], [-2.],
       [ 0.], [ 1.], [ 2.], [-1.], [-2.],
       [ 0.], [ 1.], [ 2.], [-1.], [-2.],
       [ 0.], [ 1.], [ 2.], [-1.], [-2.]], dtype=float32), 
  receivers=Array([0, 1, 2, 5, 4, 1, 2, 3, 0, 5, 2, 3, 4, 1, 0, 3, 4, 5, 2, 1, 4, 5,
       0, 3, 2, 5, 0, 1, 4, 3], dtype=int32), 
  senders=Array([0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 4, 4,
       4, 4, 4, 5, 5, 5, 5, 5], dtype=int32), 
  globals=Array([[1.]], dtype=float32), n_node=Array([6], dtype=int32), n_edge=Array([30], dtype=int32)),
GraphsTuple(nodes=Array([[ 1.0018804 ,  1.8181608 ],
       [-1.2903361 ,  0.27553296],
       [ 0.43631706,  0.3780051 ],
       [ 0.5541917 , -0.17729115],
       [ 0.3262986 , -1.8358552 ],
       [ 1.1800684 ,  1.5343047 ]], dtype=float32), 
  edges=Array([[ 0.], [ 1.], [ 2.], [-1.], [-2.],
       [ 0.], [ 1.], [ 2.], [-1.], [-2.],
       [ 0.], [ 1.], [ 2.], [-1.], [-2.],
       [ 0.], [ 1.], [ 2.], [-1.], [-2.],
       [ 0.], [ 1.], [ 2.], [-1.], [-2.],
       [ 0.], [ 1.], [ 2.], [-1.], [-2.]], dtype=float32), 
  receivers=Array([0, 1, 2, 5, 4, 1, 2, 3, 0, 5, 2, 3, 4, 1, 0, 3, 4, 5, 2, 1, 4, 5,
       0, 3, 2, 5, 0, 1, 4, 3], dtype=int32), 
  senders=Array([0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 4, 4,
       4, 4, 4, 5, 5, 5, 5, 5], dtype=int32), 
  globals=Array([[1.]], dtype=float32), n_node=Array([6], dtype=int32), n_edge=Array([30], dtype=int32))]

When I try to run this code, I get the following error:
ValueError: vmap in_axes specification must be a tree prefix of the corresponding value, got specification (GraphsTuple(nodes=1, edges=1, receivers=1, senders=1, globals=1, n_node=None, n_edge=None), GraphsTuple(nodes=1, edges=1, receivers=1, senders=1, globals=1, n_node=None, n_edge=None)) for value tree PyTreeDef(([CustomNode(namedtuple[GraphsTuple], [*, *, *, *, *, *, *]), CustomNode(namedtuple[GraphsTuple], [*, *, *, *, *, *, *])], [CustomNode(namedtuple[GraphsTuple], [*, *, *, *, *, *, *]), CustomNode(namedtuple[GraphsTuple], [*, *, *, *, *, *, *])])).

Is there a way to either use in_axes or reformat my batch data structure such that I can use jax.vmap() to batch over the windows? I'd also appreciate any other tips for batching/using vmap :)

For more context, here are my jax/jraph versions:

jax                      0.4.20
jaxlib                   0.4.20
jraph                    0.0.6.dev0

Thanks,
Mia

[miamirabelli]

Fixed by using jnp.stack to stack my GraphsTuples! Updating in case anyone runs into a similar issue